diff --git a/build.zig b/build.zig index 08765bd1df..481ad34d26 100644 --- a/build.zig +++ b/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); } diff --git a/tools/zig/bolt_ggml.py b/tools/zig/bolt_ggml.py new file mode 100644 index 0000000000..612e9440cb --- /dev/null +++ b/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(" 3 else None) + print(w.shape, f"mean {w.mean():.6g} std {w.std():.6g} absmax {np.abs(w).max():.6g}") diff --git a/tools/zig/bolt_ggml_fixture.py b/tools/zig/bolt_ggml_fixture.py new file mode 100644 index 0000000000..56ec1af8ca --- /dev/null +++ b/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]) diff --git a/tools/zig/bolt_image_fixture.py b/tools/zig/bolt_image_fixture.py new file mode 100644 index 0000000000..5f33408adf --- /dev/null +++ b/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(" f32 [rows, blocks * 256]; bit-identical to bolt_ggml.iq2_xxs.""" + grid, sign = _iq2_tables() + rows = raw.shape[0] + b = raw.reshape(-1, 66) + d = b[:, :2].contiguous().view(torch.float16).float() # [n, 1] + qs = b[:, 2:].contiguous() # 8 sub-blocks of (4 grid indices, u32 aux) + pair = qs.view(-1, 8, 8) + idx = pair[:, :, :4].long() # [n, 8, 4] + aux = pair[:, :, 4:].contiguous().view(torch.int32).long().squeeze(-1) & 0xFFFFFFFF # [n, 8] + db = d * (0.5 + (aux >> 28).float()) * 0.25 # [n, 8] + shifts = torch.arange(4) * 7 + sidx = (aux[..., None] >> shifts) & 127 # [n, 8, 4] + out = grid[idx] * sign[sidx] * db[..., None, None] # [n, 8, 4, 8] + return out.reshape(rows, -1) + + +def q8_0(raw): + """U8 [rows, blocks, 34] -> f32 [rows, blocks * 32].""" + rows = raw.shape[0] + b = raw.reshape(-1, 34) + d = b[:, :2].contiguous().view(torch.float16).float() + q = b[:, 2:].contiguous().view(torch.int8).float() + return (q * d).reshape(rows, -1) + + +_E2M1 = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], dtype=torch.float32) + + +def nvfp4(packed, scale, global_scale): + """ModelOpt NVFP4: U8 [N, K/2] (column 2j low nibble, 2j+1 high), e4m3 [N, K/16], f32 [] -> f32 [N, K].""" + n = packed.shape[0] + codes = torch.stack([packed & 0xF, packed >> 4], dim=-1).reshape(n, -1).long() + w = _E2M1[codes].view(n, -1, 16) * scale.float()[..., None] + return (w * global_scale.float()).reshape(n, -1) + + +class Checkpoint: + def __init__(self, model): + self.dir = Path(model) + self.map = json.loads((self.dir / "model.safetensors.index.json").read_text())["weight_map"] + cfg = json.loads((self.dir / "config.json").read_text()) + self.algo = {k: v["quant_algo"] for k, v in cfg["quantization_config"]["quantized_layers"].items()} + self.cfg = cfg["llm_config"] + self.files = {} + + def raw(self, name): + f = self.map[name] + if f not in self.files: + self.files[f] = safe_open(str(self.dir / f), framework="pt") + return self.files[f].get_tensor(name) + + def has(self, name): + return name in self.map + + def dense(self, name): + """A Linear weight [out, in] in f32, by the format the export records for it.""" + algo = self.algo.get(name) + if algo == "IQ2_XXS": + return iq2_xxs(self.raw(name + ".weight")) + if algo == "Q8_0": + return q8_0(self.raw(name + ".weight")) + if algo == "W4A16_NVFP4": + return nvfp4(self.raw(name + ".weight"), self.raw(name + ".weight_scale"), self.raw(name + ".weight_scale_2")) + if algo is None: + return self.raw(name + ".weight").float() + raise SystemExit(f"{name}: unsupported {algo}") + + def vec(self, name): + return self.raw(name).float() + + +# ---------------------------------------------------------------- blocks + + +def rmsnorm(x, w, eps): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps) * w + + +class State: + """Per-layer recurrent state: Mamba conv window + ssm state, attention K/V.""" + + def __init__(self): + self.conv, self.ssm, self.k, self.v = {}, {}, {}, {} + self.pos = 0 + + +def mamba(ck, cfg, li, x, st): + p = f"{PRE}layers.{li}.mixer." + H, P, N, G = cfg["mamba_num_heads"], cfg["mamba_head_dim"], cfg["ssm_state_size"], cfg["n_groups"] + inner = H * P + T = x.shape[0] + zxbcdt = x @ ck.dense(p + "in_proj").T # [T, inner + conv_dim + H] + conv_dim = inner + 2 * G * N + z, xbc, dt = zxbcdt.split([inner, conv_dim, H], dim=-1) + + # causal depthwise conv over [previous K-1 inputs, these T], + bias, SiLU + cw = ck.vec(p + "conv1d.weight")[:, 0, :] # [conv_dim, K] + K = cw.shape[1] + prev = st.conv.get(li, torch.zeros(K - 1, conv_dim)) + seq = torch.cat([prev, xbc], 0) # [K-1+T, conv_dim] + st.conv[li] = seq[-(K - 1):].clone() + win = seq.unfold(0, K, 1) # [T, conv_dim, K] + xbc = F.silu((win * cw).sum(-1) + ck.vec(p + "conv1d.bias")) + xs, B, C = xbc.split([inner, G * N, G * N], dim=-1) + + # dt = softplus(dt + dt_bias), limited to time_step_limit = (0, inf): no clamp (the kernel path) + dt = F.softplus(dt + ck.vec(p + "dt_bias")) # [T, H] + lo, hi = cfg["time_step_limit"][0], cfg["time_step_limit"][1] + hi = float(hi["__float__"]) if isinstance(hi, dict) else float(hi) + if lo > 0 or hi != math.inf: + dt = dt.clamp(lo, hi) + A = -torch.exp(ck.vec(p + "A_log")) # [H] + D = ck.vec(p + "D") + xs = xs.view(T, H, P) + rep = H // G + B = B.view(T, G, N).repeat_interleave(rep, dim=1) # [T, H, N] + C = C.view(T, G, N).repeat_interleave(rep, dim=1) + dA = torch.exp(dt * A) # [T, H] + h = st.ssm.get(li, torch.zeros(H, P, N)) + ys = torch.empty(T, H, P) + for t in range(T): + h = h * dA[t][:, None, None] + (dt[t][:, None] * xs[t])[:, :, None] * B[t][:, None, :] + ys[t] = torch.einsum("hpn,hn->hp", h, C[t]) + st.ssm[li] = h + y = (ys + xs * D[:, None]).reshape(T, inner) + + # gated RMSNorm, norm_before_gate=False: normalize y * silu(z) per group of inner / n_groups + g = (y * F.silu(z)).view(T, G, inner // G) + g = g * torch.rsqrt(g.pow(2).mean(-1, keepdim=True) + cfg["layer_norm_epsilon"]) + y = g.reshape(T, inner) * ck.vec(p + "norm.weight") + return y @ ck.dense(p + "out_proj").T + + +def attention(ck, cfg, li, x, st): + p = f"{PRE}layers.{li}.mixer." + nh, nkv, hd = cfg["num_attention_heads"], cfg["num_key_value_heads"], cfg["head_dim"] + T = x.shape[0] + q = (x @ ck.dense(p + "q_proj").T).view(T, nh, hd) + k = (x @ ck.dense(p + "k_proj").T).view(T, nkv, hd) + v = (x @ ck.dense(p + "v_proj").T).view(T, nkv, hd) + if li in st.k: + k = torch.cat([st.k[li], k], 0) + v = torch.cat([st.v[li], v], 0) + st.k[li], st.v[li] = k, v + S = k.shape[0] + rep = nh // nkv + kk = k.repeat_interleave(rep, dim=1).transpose(0, 1) # [nh, S, hd] + vv = v.repeat_interleave(rep, dim=1).transpose(0, 1) + s = torch.einsum("thd,hsd->hts", q, kk) * hd**-0.5 + qpos = torch.arange(S - T, S)[:, None] + s = s.masked_fill(torch.arange(S)[None, :] > qpos, float("-inf")) + o = torch.einsum("hts,hsd->thd", torch.softmax(s, -1), vv).reshape(T, nh * hd) + return o @ ck.dense(p + "o_proj").T + + +def moe(ck, cfg, bc, li, x, routes): + p = f"{PRE}layers.{li}.mixer." + E, k = bc["n_routed_experts"], bc["num_experts_per_tok"] + T = x.shape[0] + # router: sigmoid(x W^T); choice by score + bias (n_group 1); unbiased weights, normalized, times the scaling factor + assert cfg["n_group"] == 1 and cfg["topk_group"] == 1 + scores = torch.sigmoid(x @ ck.vec(p + "gate.weight").T) # [T, E] + choice = scores + ck.vec(p + "gate.e_score_correction_bias") + idx = torch.topk(choice, k, dim=-1).indices # [T, k], by descending choice + w = scores.gather(1, idx) + if cfg["norm_topk_prob"]: + w = w / (w.sum(-1, keepdim=True) + 1e-20) + w = w * cfg["routed_scaling_factor"] + routes[li] = (idx.clone(), w.clone()) + + lat = x @ ck.dense(p + "fc1_latent_proj").T # [T, latent] + acc = torch.zeros_like(lat) + for e in torch.unique(idx).tolist(): + tok, slot = torch.where(idx == e) + up = ck.dense(f"{p}experts.{e}.up_proj") # [inter, latent] + down = ck.dense(f"{p}experts.{e}.down_proj") # [latent, inter] + hdn = F.relu(lat[tok] @ up.T).square() + acc.index_add_(0, tok, (hdn @ down.T) * w[tok, slot][:, None]) + routed = acc @ ck.dense(p + "fc2_latent_proj").T + sh = F.relu(x @ ck.dense(p + "shared_experts.up_proj").T).square() + return routed + sh @ ck.dense(p + "shared_experts.down_proj").T + + +def lm_head(ck, x, rows=16384): + raw = ck.raw("language_model.lm_head.weight") + return torch.cat([x @ q8_0(raw[i:i + rows]).T for i in range(0, raw.shape[0], rows)], -1) + + +def forward(ck, ids, st=None, dump=None, dump_last=8, logits_last=None, verbose=True, splice=None): + """Run `ids` (continuing `st`) -> (logits, state, routes); `splice` = (positions, rows) overwrites embeddings.""" + cfg = ck.cfg + st = st or State() + eps = cfg["layer_norm_epsilon"] + ids_t = torch.tensor(ids, dtype=torch.long) + emb = ck.raw(PRE + "embeddings.weight") + h = emb[ids_t].float() + if splice is not None: + pos, rows = splice + h[pos] = rows.float() + routes = {} + keep = slice(-dump_last, None) if dump_last else slice(None) + if dump: + np.save(dump / "h_embed.npy", h[keep].numpy()) + t0 = time.time() + for li, bc in enumerate(cfg["block_configs"]): + t1 = time.time() + x = rmsnorm(h, ck.vec(f"{PRE}layers.{li}.norm.weight"), eps) + kind = bc["block_type"] + if kind == "mamba": + out = mamba(ck, cfg, li, x, st) + elif kind == "attention": + out = attention(ck, cfg, li, x, st) + elif kind == "moe": + out = moe(ck, cfg, bc, li, x, routes) + else: + raise SystemExit(f"layer {li}: {kind}") + h = h + out + if dump: + np.save(dump / f"h_layer{li:02d}.npy", h[keep].numpy()) + np.save(dump / f"mixer_layer{li:02d}.npy", out[keep].numpy()) + if li in routes: + np.save(dump / f"route_idx_layer{li:02d}.npy", routes[li][0][keep].numpy()) + np.save(dump / f"route_w_layer{li:02d}.npy", routes[li][1][keep].numpy()) + if verbose: + print(f" layer {li:2d} {kind:9s} |h|={h[-1].norm():9.3f} {time.time() - t1:5.1f}s", file=sys.stderr, flush=True) + st.pos += len(ids) + h = rmsnorm(h, ck.vec(PRE + "norm_f.weight"), eps) + if logits_last: + h = h[-logits_last:] + logits = lm_head(ck, h) + if dump: + np.save(dump / "h_final_norm.npy", h[keep].numpy()) + np.save(dump / "logits.npy", logits[keep].numpy()) + if verbose: + print(f" forward {len(ids)} tokens: {time.time() - t0:.1f}s", file=sys.stderr) + return logits, st, routes + + +# ---------------------------------------------------------------- fixtures / scoring + + +def load_fixture(name, path=FIXTURES): + for line in open(path): + d = json.loads(line) + if d["name"] == name: + return d + raise SystemExit(f"no fixture {name!r} in {path}") + + +def score(lp, tops, label): + """lp [n, vocab] log-probs vs vLLM dicts (token id -> logprob) for the same n positions.""" + strict = tie = 0 + diffs = [] + rows = [] + for i, top in enumerate(tops): + ids = torch.tensor([int(t) for t in top]) + ref = torch.tensor(list(top.values())) + mine = int(lp[i].argmax()) + best = ref.max() + strict += mine == int(ids[ref.argmax()]) + tie += bool((ids[ref >= best - 1e-6] == mine).any()) + err = (lp[i][ids] - ref).abs() + d = err.max().item() + diffs.append(d) + j = int(err.argmax()) + rows.append((d, i, int(ids[j]), ref[j].item(), lp[i][ids[j]].item())) + n = len(tops) + diffs = np.array(diffs) + print(f"{label}: top-1 {strict}/{n} = {100 * strict / n:.1f}% (counting vLLM ties: {tie}/{n} = {100 * tie / n:.1f}%) | " + f"max|dlogprob| over vLLM top-20: mean {diffs.mean():.3f} median {np.median(diffs):.3f} max {diffs.max():.3f}") + for d, i, t, r, m in sorted(rows, reverse=True)[:3]: + print(f" worst: position {i} id {t}: vLLM {r:.3f} ours {m:.3f} (|d| {d:.3f})") + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + src = ap.add_mutually_exclusive_group(required=True) + src.add_argument("--ids", help="file of whitespace/comma separated token ids (or a JSON list)") + src.add_argument("--fixture", help="fixture name in text.jsonl (short, code, think, multiturn, long, german)") + ap.add_argument("--model", default=MODEL) + ap.add_argument("--fixtures", default=FIXTURES) + ap.add_argument("--dump", help="directory for per-layer npy dumps") + ap.add_argument("--dump-last", type=int, default=8, help="positions kept in dumps (0 = all)") + ap.add_argument("--compare", action="store_true", help="score prompt logprobs against the fixture") + ap.add_argument("--generate", type=int, default=0, help="greedy tokens to generate after the prompt") + ap.add_argument("--threads", type=int, default=16) + ap.add_argument("--quiet", action="store_true") + a = ap.parse_args() + torch.set_num_threads(a.threads) + torch.set_grad_enabled(False) + ck = Checkpoint(a.model) + + fx = None + if a.fixture: + fx = load_fixture(a.fixture, a.fixtures) + ids = fx["prompt_ids"] + else: + text = Path(a.ids).read_text().strip() + ids = json.loads(text) if text.startswith("[") else [int(t) for t in text.replace(",", " ").split()] + dump = None + if a.dump: + dump = Path(a.dump) + dump.mkdir(parents=True, exist_ok=True) + (dump / "ids.json").write_text(json.dumps(ids)) + + t0 = time.time() + logits, st, _ = forward(ck, ids, dump=dump, dump_last=a.dump_last, verbose=not a.quiet) + print(f"prompt {len(ids)} tokens in {time.time() - t0:.1f}s", flush=True) + lp = torch.log_softmax(logits, -1) + if a.compare: + if fx is None: + raise SystemExit("--compare needs --fixture") + # prompt_top[i] is vLLM's distribution for position i given ids[:i], i.e. our logits at i-1 + score(lp[:-1], fx["prompt_top"][1:], f"{a.fixture} prompt ({len(ids) - 1} positions)") + + if a.generate: + from tokenizers import Tokenizer + tok = Tokenizer.from_file(str(Path(a.model) / "tokenizer.json")) + out = [int(logits[-1].argmax())] + steps = [lp[-1]] + for _ in range(a.generate - 1): + lg, st, _ = forward(ck, [out[-1]], st, verbose=False) + steps.append(torch.log_softmax(lg[-1], -1)) + out.append(int(lg[-1].argmax())) + print(f"generated {len(out)} tokens in {time.time() - t0:.1f}s total") + print("ids:", out) + print("text:", repr(tok.decode(out))) + if fx is not None: + ref = fx["gen_tokens"][: len(out)] + mine = [tok.decode([i]) for i in out] + match = 0 + for m, r in zip(mine, ref): + if m != r: + break + match += 1 + print(f"greedy prefix match vs vLLM: {match}/{len(ref)}") + print("vLLM:", repr("".join(ref))) + # logprob diffs at the generated positions (vLLM keys are token strings) + diffs = [] + for s_lp, top in zip(steps[:match + 1], fx["gen_top"]): + d = 0.0 + for i in torch.topk(s_lp, 40).indices.tolist(): + t = tok.decode([i]) + if t in top: + d = max(d, abs(s_lp[i].item() - top[t])) + diffs.append(d) + print(f"gen max|dlogprob| (shared top ids, first {len(diffs)} steps): mean {np.mean(diffs):.3f} max {np.max(diffs):.3f}") + + +if __name__ == "__main__": + main() diff --git a/tools/zig/bolt_vision_reference.py b/tools/zig/bolt_vision_reference.py new file mode 100644 index 0000000000..361dde882b --- /dev/null +++ b/tools/zig/bolt_vision_reference.py @@ -0,0 +1,296 @@ +"""fp32 PyTorch reference of Bolt's image path (preprocess, C-RADIO tower, projector), spliced into the LM one.""" + +import argparse +import copy +import json +import math +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import bolt_reference as lm # noqa: E402 + +FIXTURES = "/storage/models/bolt-fixtures/vision.jsonl" +VPRE = "vision_model.radio_model." + + +# ---------------------------------------------------------------- preprocessing + + +def target_patches(w, h, cfg): + """(w_patches, h_patches) of the single dynamic-resolution tile (image path, one image).""" + ps, mn, mx = cfg["patch_size"], cfg["min_num_patches"], cfg["max_num_patches"] + r = int(round(1 / cfg["downsample_ratio"])) + budget = (cfg["max_model_len"] - 4) * r * r # vLLM subtracts the text length too; capped anyway + avail = max(min(budget, mx), mn) + ph = round(h / ps + 0.5) + pw = round(w / ps + 0.5) + f = min(math.sqrt(avail / (ph * pw)), 1.0) + th, tw = math.floor(f * ph), math.floor(f * pw) + if avail > mn and th * tw < mn: + up = math.sqrt(mn / (th * tw)) + th, tw = math.ceil(up * th), math.ceil(up * tw) + if th % r: + th = th + (r - th % r) if (th + r - th % r) * tw <= avail else max(r, th - th % r) + if tw % r: + tw = tw + (r - tw % r) if th * (tw + r - tw % r) <= avail else max(r, tw - tw % r) + return tw, th + + +def preprocess(path, cfg): + img = Image.open(path) + if img.mode != "RGB": + img = img.convert("RGB") + wp, hp = target_patches(img.width, img.height, cfg) + ps = cfg["patch_size"] + t = torch.from_numpy(np.array(img, dtype=np.uint8)).permute(2, 0, 1)[None].float() + if t.shape[-2:] != (hp * ps, wp * ps): + t = F.interpolate(t, size=(hp * ps, wp * ps), mode="bicubic", align_corners=False, antialias=True) + mean = torch.tensor(cfg["norm_mean"]).view(1, 3, 1, 1) + std = torch.tensor(cfg["norm_std"]).view(1, 3, 1, 1) + return ((t / 255.0 - mean) / std)[0], (img.width, img.height), (hp, wp) + + +# ---------------------------------------------------------------- C-RADIO tower + + +def pos_embed(table, hp, wp): + """CPE: [1, R*R, C] table -> [hp*wp, C]; bilinear (align_corners=False) to max side, crop.""" + n = int(math.isqrt(table.shape[1])) + pos = table.reshape(1, n, n, -1).permute(0, 3, 1, 2).float() + m = max(hp, wp) + pos = F.interpolate(pos, size=(m, m), mode="bilinear", align_corners=False)[..., :hp, :wp] + return pos.flatten(2)[0].T + + +def radio(ck, vcfg, px, dump=None, dump_layers=False, verbose=True): + """px [3, H, W] normalized -> (all tokens [10 + hp*wp, 1280], patch features [hp*wp, 1280]).""" + ps, C, nh = vcfg["patch_size"], vcfg["hidden_size"], vcfg["num_attention_heads"] + eps = vcfg["layer_norm_eps"] + _, H, W = px.shape + hp, wp = H // ps, W // ps + m = VPRE + "model." + g = lambda n: ck.raw(n).float() # noqa: E731 + patches = px.reshape(3, hp, ps, wp, ps).permute(1, 3, 0, 2, 4).reshape(hp * wp, 3 * ps * ps) + x = patches @ g(m + "patch_generator.embedder.weight").T + x = x + pos_embed(g(m + "patch_generator.pos_embed"), hp, wp) + prefix = g(m + "patch_generator.cls_token.token") # [3 cls + 7 registers, C] + x = torch.cat([prefix, x], 0) + if dump: + np.save(dump / "tower_embed.npy", x.numpy()) + hd = C // nh + T = x.shape[0] + for i in range(vcfg["num_hidden_layers"]): + t0 = time.time() + b = f"{m}blocks.{i}." + y = F.layer_norm(x, (C,), g(b + "norm1.weight"), g(b + "norm1.bias"), eps) + qkv = (y @ g(b + "attn.qkv.weight").T + g(b + "attn.qkv.bias")).view(T, 3, nh, hd).permute(1, 2, 0, 3) + o = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2]) # [nh, T, hd], scale hd^-0.5 + o = o.transpose(0, 1).reshape(T, C) + x = x + o @ g(b + "attn.proj.weight").T + g(b + "attn.proj.bias") + y = F.layer_norm(x, (C,), g(b + "norm2.weight"), g(b + "norm2.bias"), eps) + y = F.gelu(y @ g(b + "mlp.fc1.weight").T + g(b + "mlp.fc1.bias")) + x = x + y @ g(b + "mlp.fc2.weight").T + g(b + "mlp.fc2.bias") + if dump and dump_layers: + np.save(dump / f"tower_layer{i:02d}.npy", x.numpy()) + if verbose: + print(f" radio block {i:2d} |x|={x[10:].norm(dim=-1).mean():8.3f} {time.time() - t0:5.1f}s", + file=sys.stderr, flush=True) + skip = vcfg["num_cls_tokens"] + vcfg["num_registers"] + return x, x[skip:] + + +def pixel_shuffle(feat, hp, wp, r=2): + """ps_version v2: [hp*wp, C] -> [(hp/r)*(wp/r), r*r*C], channel index (dy*r + dx)*C + c.""" + C = feat.shape[-1] + x = feat.view(hp // r, r, wp // r, r, C).permute(0, 2, 1, 3, 4) + return x.reshape((hp // r) * (wp // r), r * r * C) + + +def projector(ck, x): + w0 = ck.raw("mlp1.0.weight").float() + x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5) * w0 + x = F.relu(x @ ck.raw("mlp1.1.weight").float().T).square() + return x @ ck.raw("mlp1.3.weight").float().T + + +# ---------------------------------------------------------------- prompt + + +def build_prompt(question, n_img, image_first=True): + """The chat template (no system content, thinking off) with the processor's image expansion.""" + img = "" + "" * n_img + "" + content = img + question if image_first else question + img + return ("<|im_start|>system\n<|im_end|>\n<|im_start|>user\n" + content + + "<|im_end|>\n<|im_start|>assistant\n") + + +# ---------------------------------------------------------------- scoring + + +def compare_step(tok, lp, top, cache): + """lp [vocab] vs vLLM {token string: logprob}; returns (n mapped, max |d|, worst).""" + if "rev" not in cache or cache["rev_lp"] is not lp: + rev = {} + for i in torch.topk(lp, 4000).indices.tolist(): + rev.setdefault(tok.decode([i]), i) + cache["rev"], cache["rev_lp"] = rev, lp + rev = cache["rev"] + d, worst, n = 0.0, None, 0 + for s, ref in top.items(): + i = rev.get(s) + if i is None: + ids = tok.encode(s, add_special_tokens=False).ids + if len(ids) != 1: + continue + i = ids[0] + n += 1 + e = abs(lp[i].item() - ref) + if e > d: + d, worst = e, (s, ref, lp[i].item()) + return n, d, worst + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--case", type=int, default=0, help="line index in vision.jsonl") + ap.add_argument("--image", help="image file (overrides the fixture's; needs --question)") + ap.add_argument("--question") + ap.add_argument("--model", default=lm.MODEL) + ap.add_argument("--fixtures", default=FIXTURES) + ap.add_argument("--text-first", action="store_true", help="question before the image in the user turn") + ap.add_argument("--vision-only", action="store_true", help="stop after the projector") + ap.add_argument("--generate", type=int, default=8, help="free-running greedy tokens") + ap.add_argument("--forced", type=int, default=16, help="teacher-forced steps scored against the fixture") + ap.add_argument("--dump", help="directory for npy dumps (pixels, tower, shuffle, projector)") + ap.add_argument("--dump-tower-layers", action="store_true", help="also every RADIO block output") + ap.add_argument("--dump-lm", action="store_true", help="also the LM per-layer dumps (last 8 positions)") + ap.add_argument("--threads", type=int, default=16) + a = ap.parse_args() + torch.set_num_threads(a.threads) + torch.set_grad_enabled(False) + + root = json.loads((Path(a.model) / "config.json").read_text()) + pcfg = json.loads((Path(a.model) / "preprocessor_config.json").read_text()) + vcfg = root["vision_config"] + assert root["ps_version"] == "v2" and not root["use_thumbnail"] and not root.get("vision_final_layernorm") + assert not vcfg["use_swiglu_ffn"] and vcfg["hidden_act"] == "gelu" + + fx = None + if a.image: + img_path, question = a.image, a.question + else: + fx = json.loads(open(a.fixtures).readlines()[a.case]) + img_path, question = str(Path(a.fixtures).parent / fx["image"]), fx["question"] + + dump = None + if a.dump: + dump = Path(a.dump) + dump.mkdir(parents=True, exist_ok=True) + + t0 = time.time() + px, (w, h), (hp, wp) = preprocess(img_path, pcfg) + n_img = hp * wp // 4 + print(f"image {Path(img_path).name} {w}x{h} -> {wp * 16}x{hp * 16} px, patch grid {hp}x{wp} = {hp * wp}, " + f"{n_img} image tokens", flush=True) + + from tokenizers import Tokenizer + tok = Tokenizer.from_file(str(Path(a.model) / "tokenizer.json")) + ids = tok.encode(build_prompt(question, n_img, not a.text_first), add_special_tokens=False).ids + img_pos = torch.tensor([i for i, t in enumerate(ids) if t == root["img_context_token_id"]]) + assert len(img_pos) == n_img + line = f"prompt tokens {len(ids)}" + if fx: + line += f" (vLLM usage.prompt_tokens {fx['usage']['prompt_tokens']}: " \ + f"{'MATCH' if len(ids) == fx['usage']['prompt_tokens'] else 'MISMATCH'})" + print(line, flush=True) + + ck = lm.Checkpoint(a.model) + allx, feat = radio(ck, vcfg, px, dump, a.dump_tower_layers) + shuf = pixel_shuffle(feat, hp, wp, int(round(1 / root["downsample_ratio"]))) + proj = projector(ck, shuf) + print(f"vision: tower {tuple(allx.shape)} -> features {tuple(feat.shape)} -> shuffle {tuple(shuf.shape)} " + f"-> projector {tuple(proj.shape)} |row| mean {proj.norm(dim=-1).mean():.3f} {time.time() - t0:.1f}s", + flush=True) + if dump: + np.save(dump / "pixel_values.npy", px.numpy()) + np.save(dump / "image_rgb.npy", np.array(Image.open(img_path).convert("RGB"), dtype=np.uint8)) + np.save(dump / "tower_out.npy", allx.numpy()) + np.save(dump / "tower_features.npy", feat.numpy()) + np.save(dump / "pixel_shuffle.npy", shuf.numpy()) + np.save(dump / "projector_out.npy", proj.numpy()) + (dump / "ids.json").write_text(json.dumps(ids)) + (dump / "meta.json").write_text(json.dumps({"image": str(img_path), "orig_wh": [w, h], + "grid_hw": [hp, wp], "n_image_tokens": n_img, + "image_positions": [int(img_pos[0]), int(img_pos[-1])]})) + if a.vision_only: + return + + t1 = time.time() + logits, st, _ = lm.forward(ck, ids, splice=(img_pos, proj), dump=dump if a.dump_lm else None, logits_last=1) + print(f"LM prompt {len(ids)} tokens in {time.time() - t1:.1f}s", flush=True) + lp = torch.log_softmax(logits[-1], -1) + st0 = copy.deepcopy(st) + out = [int(lp.argmax())] + steps = [lp] + for _ in range(a.generate - 1): + lg, st, _ = lm.forward(ck, [out[-1]], st, verbose=False) + steps.append(torch.log_softmax(lg[-1], -1)) + out.append(int(lg[-1].argmax())) + mine = [tok.decode([i]) for i in out] + print("ours:", repr("".join(mine)), out) + if fx is None: + print(f"total {time.time() - t0:.1f}s") + return + ref = [g["t"] for g in fx["gen"]] + print("vLLM:", repr("".join(ref[:len(out)]))) + match = 0 + for m_, r_ in zip(mine, ref): + if m_ != r_: + break + match += 1 + print(f"free-running greedy prefix match vs vLLM: {match}/{len(out)}") + top = torch.topk(lp, 5) + print(" step 0 ours top-5:", [(tok.decode([i]), round(v, 3)) for v, i in zip(top.values.tolist(), top.indices.tolist())]) + rank = sorted(fx["gen"][0]["top"], key=fx["gen"][0]["top"].get, reverse=True) + print(" step 0 vLLM top-5:", [(r, round(fx["gen"][0]["top"][r], 3)) for r in rank[:5]]) + + # teacher-forced: feed vLLM's tokens, score every step (a vLLM near-tie can't derail the rest) + n = min(max(a.forced, 1), len(fx["gen"])) + st, s_lp, cache = st0, lp, {} + strict = tie = 0 + diffs = [] + for k in range(n): + g = fx["gen"][k] + cnt, d, worst = compare_step(tok, s_lp, g["top"], cache) + o = int(s_lp.argmax()) + os_ = tok.decode([o]) + best = max(g["top"].values()) + is_tie = os_ in g["top"] and g["top"][os_] >= best - 1e-4 + strict += os_ == g["t"] + tie += is_tie + diffs.append(d) + rank = sorted(g["top"], key=g["top"].get, reverse=True) + print(f" forced step {k:2d}: ours {os_!r} {s_lp[o].item():7.3f} | vLLM {g['t']!r} {g['lp']:7.3f} " + f"(2nd {rank[1]!r} {g['top'][rank[1]]:7.3f}) | mapped {cnt}/20 max|d| {d:.3f}" + + ("" if os_ == g["t"] else f" {'TIE' if is_tie else 'MISS'} worst {worst}")) + if k + 1 < n: + nxt = cache["rev"].get(g["t"]) + if nxt is None: + nxt = tok.encode(g["t"], add_special_tokens=False).ids[0] + lg, st, _ = lm.forward(ck, [nxt], st, verbose=False) + s_lp = torch.log_softmax(lg[-1], -1) + diffs = np.array(diffs) + print(f"teacher-forced {n} steps: top-1 {strict}/{n} (counting vLLM ties {tie}/{n}) | max|dlogprob| over " + f"vLLM top-20: mean {diffs.mean():.3f} median {np.median(diffs):.3f} max {diffs.max():.3f}") + print(f"total {time.time() - t0:.1f}s") + + +if __name__ == "__main__": + main() diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..004e23f056 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 vision_build = @import("vision.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 }; @@ -32,6 +33,7 @@ const kernels = [_]Kernel{ .{ .name = "lane_gemv", .flags = &.{"-O3"} }, // ours: qmm_group's arithmetic, a column tile's K slices in one CTA .{ .name = "sample", .flags = &.{ "-O3", "--fmad=false", "--ftz=false" } }, // ours: the Metal engine's keyed draws .{ .name = "train", .flags = &.{"-O3"} }, // ours: Sliding Weights' learning step (train.metal's kernels) + .{ .name = "radio_vision", .flags = &.{ "-O3", "--fmad=false" } }, // ours: Nemotron 3.5 Bolt's C-RADIO tower and projector .{ .name = "nemotron_norms", .flags = glue }, // ours, each with a host reference (glue_ref.zig): Nemotron's glue .{ .name = "nemotron_route", .flags = glue }, .{ .name = "nemotron_mamba", .flags = glue }, @@ -42,6 +44,7 @@ const kernels = [_]Kernel{ .{ .name = "qmmf", .flags = &.{"-O3"} }, // nvfp4/qmmf.cu's device code (FP8G, FP4), tensorfold_nvfp4_v3 .{ .name = "prompt16", .flags = &.{"-O3"} }, // nvfp4/prompt.cu's FP4 and FP8 tile 4, tensorfold_nvfp4_prompt_v1 .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts + .{ .name = "ggml", .flags = &.{"-O3"} }, // ours: Bolt's IQ2_XXS and Q8_0 blocks, grouped experts and dense rows .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, @@ -129,6 +132,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu const runner = b.createModule(.{ .root_source_file = b.path("zig/tests/cuda/main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); runner.addImport("cuda", cuda); runner.addImport("core", mods.core); + runner.addImport("qwen_image", vision_build.module(b, target, optimize)); runner.addImport("lanes", mods.lanes); runner.addImport("nemotron", mods.nemotron); const exlr8 = exlr8_build.coreAlias(b, target, optimize, mods.core); @@ -139,12 +143,18 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); + // `zig build bolt-vision`: Bolt's image tower against the reference's dumps (zig/tests/bolt_vision_main.zig) + const bolt = b.createModule(.{ .root_source_file = b.path("zig/tests/bolt_vision_main.zig"), .target = target, .optimize = optimize, .link_libc = true, .strip = strip }); + bolt.addImport("cuda", cuda); + bolt.addImport("nemotron", mods.nemotron); + bolt.addImport("qwen_image", vision_build.module(b, target, optimize)); + b.step("bolt-vision", "tf-bolt-vision: Bolt's image tower and projector against bolt_vision_reference.py's dumps").dependOn(&b.addInstallArtifact(b.addExecutable(.{ .name = "tf-bolt-vision", .root_module = bolt }), .{}).step); nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { - 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 = 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 = lanes }, .{ .name = "qwen_image", .module = vision_build.module(b, target, optimize) } } }); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, @@ -210,7 +220,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug), vision_build.module(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); diff --git a/zig/build/vision.zig b/zig/build/vision.zig new file mode 100644 index 0000000000..0b78064367 --- /dev/null +++ b/zig/build/vision.zig @@ -0,0 +1,6 @@ +//! The image-input module (vision/qwen_image.zig), shared by the engine API and the CUDA test runner. +const std = @import("std"); + +pub fn module(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + return b.createModule(.{ .root_source_file = b.path("zig/src/vision/qwen_image.zig"), .target = target, .optimize = optimize, .link_libc = true }); +} diff --git a/zig/kernels/cuda/ggml.cu b/zig/kernels/cuda/ggml.cu new file mode 100644 index 0000000000..10d54dd578 --- /dev/null +++ b/zig/kernels/cuda/ggml.cu @@ -0,0 +1,456 @@ +// Bolt's ggml blocks (IQ2_XXS, Q8_0) on mma.sync, weights as A and rows as B: grouped experts and dense rows. + +#include +#include +#include + +#include "experts.cuh" +#include "ggml_iq2xxs_tables.cuh" + +// Lane t of a 32-input group holds inputs 8t .. 8t + 3 (step 0) and 8t + 4 .. 8t + 7 (step 1), in both operands. +namespace tf_ggml { + +__device__ __forceinline__ uint4 ld_nc(const uint8_t* p) { + uint4 r; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0, %1, %2, %3}, [%4];\n" + : "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w) + : "l"(p)); + return r; +} + +__device__ __forceinline__ uint32_t ld_nc32(const uint8_t* p) { + uint32_t r; + asm volatile("ld.global.nc.L1::no_allocate.u32 %0, [%1];\n" : "=r"(r) : "l"(p)); + return r; +} + +__device__ __forceinline__ float half_f(uint32_t bits) { return __half2float(__ushort_as_half((unsigned short)bits)); } + +// The sign bits of a bf16 pair from bits 0 and 1 of sb (< 2^16): bit 0 lands on 15, bit 1 on 31, no other bit there. +__device__ __forceinline__ uint32_t sign_mask(uint32_t sb) { return (sb * 0x40008000u) & 0x80008000u; } + +// Two grid bytes of g (selector sel) as exact bf16 integers, signs from bits 0 and 1 of sb. +__device__ __forceinline__ uint32_t grid_pair(uint32_t g, uint32_t sel, uint32_t sb) { + uint32_t v = __byte_perm(g, 0x43u, sel); + const uint32_t k = K128; + __nv_bfloat162 r = __hsub2(*reinterpret_cast<__nv_bfloat162*>(&v), *reinterpret_cast(&k)); + return *reinterpret_cast(&r) ^ sign_mask(sb); +} + +// IQ2_XXS: a group is a 32-weight sub-block (4 grid indices, 4 x 7 sign bits, a 4-bit scale); 8 groups a block. +struct Iq2 { + static constexpr int G = 8, BYTES = 1056, Q = 4; // Q: groups a decode load brings, a lane's 16 bytes each + static constexpr bool TABLE = true; + // lane (gq, t) loads group 4q + t's bytes of columns gq, gq + 8: a warp's load is 512 distinct bytes + static __device__ __forceinline__ const uint8_t* uptr(const uint8_t* tile, int q, int lane) { + return tile + (q >> 1) * BYTES + (q & 1) * 512 + (lane & 3) * 128 + (lane >> 2) * 16; + } + static __device__ __forceinline__ const uint8_t* wptr(const uint8_t* tile, int g, int lane) { + return tile + (g >> 3) * BYTES + (g & 7) * 128 + (lane >> 2) * 16; + } + static __device__ __forceinline__ const uint8_t* dptr(const uint8_t* tile, int g, int lane) { + return tile + (g >> 3) * BYTES + 1024 + (lane >> 2) * 4; + } + static __device__ __forceinline__ void decode(const uint4& w, uint32_t dp, int t, const uint2* tab, + uint32_t (&a)[2][4], float (&sc)[2]) { + sc[0] = half_f(dp & 0xFFFFu) * (0.5f + (float)(w.y >> 28)) * 0.25f; + sc[1] = half_f(dp >> 16) * (0.5f + (float)(w.w >> 28)) * 0.25f; +#pragma unroll + for (int c = 0; c < 2; ++c) { // lane t decodes grid entry t of the group: 8 weights, both steps + const uint32_t aux0 = c ? w.z : w.x, aux1 = c ? w.w : w.y; + const uint2 g = tab[(aux0 >> (8 * t)) & 255u]; + const uint32_t si = (aux1 >> (7 * t)) & 127u; + const uint32_t sb = si | ((__popc(si) & 1u) << 7); // ksigns_iq2xs: even parity in bit 7 +#pragma unroll + for (int s = 0; s < 2; ++s) { + const uint32_t gs = s ? g.y : g.x, bs = sb >> (4 * s); + a[s][c] = grid_pair(gs, 0x4140u, bs); + a[s][2 + c] = grid_pair(gs, 0x4342u, bs >> 2); + } + } + } +}; + +// Two int8 of u (bytes 2h, 2h + 1) as exact bf16. +__device__ __forceinline__ uint32_t i8_pair(uint32_t u, int h) { + const uint32_t v = u ^ 0x80808080u; + const float f0 = __uint_as_float(__byte_perm(v, 0x4B000000u, 0x7440u | (2 * h))) - 8388736.f; + const float f1 = __uint_as_float(__byte_perm(v, 0x4B000000u, 0x7440u | (2 * h + 1))) - 8388736.f; + __nv_bfloat162 r = __floats2bfloat162_rn(f0, f1); + return *reinterpret_cast(&r); +} + +// Q8_0: a group is one block (fp16 d, 32 int8). +struct Q8 { + static constexpr int G = 1, BYTES = 544, Q = 1; + static constexpr bool TABLE = false; + static __device__ __forceinline__ const uint8_t* uptr(const uint8_t* tile, int q, int lane) { + return tile + q * BYTES + lane * 16; + } + static __device__ __forceinline__ const uint8_t* wptr(const uint8_t* tile, int g, int lane) { + return tile + g * BYTES + lane * 16; + } + static __device__ __forceinline__ const uint8_t* dptr(const uint8_t* tile, int g, int lane) { + return tile + g * BYTES + 512 + (lane >> 2) * 4; + } + static __device__ __forceinline__ void decode(const uint4& w, uint32_t dp, int, const uint2*, + uint32_t (&a)[2][4], float (&sc)[2]) { + sc[0] = half_f(dp & 0xFFFFu); + sc[1] = half_f(dp >> 16); +#pragma unroll + for (int s = 0; s < 2; ++s) +#pragma unroll + for (int c = 0; c < 2; ++c) { + const uint32_t u = comp(w, 2 * s + c); + a[s][c] = i8_pair(u, 0); + a[s][2 + c] = i8_pair(u, 1); + } + } +}; + +// IQ2_XXS's grid in shared memory (indices diverge across a warp); signs and the bf16 conversion are arithmetic. +constexpr int TABLE_WORDS = 256; + +template +__device__ __forceinline__ void load_tables(uint2* tab) { + if constexpr (F::TABLE) { + for (int i = threadIdx.x; i < TABLE_WORDS; i += blockDim.x) { + const uint64_t g = tf_iq2xxs_grid[i]; + tab[i] = make_uint2((uint32_t)g, (uint32_t)(g >> 32)); + } + __syncthreads(); + } +} + +// One item of rows: grouped (expert, first member, count) from the plan, or a dense run of rows (members null). +struct Item { + int e, first, cnt; +}; + +__device__ __forceinline__ int pair_of(const int* __restrict__ members, const Item& it, int r) { + return members ? __ldg(members + it.first + r) : it.first + r; +} + +template +__device__ __forceinline__ void store(void* out, size_t at, float v) { + if constexpr (EPI == 0) { + reinterpret_cast(out)[at] = v; + } else { + reinterpret_cast<__nv_bfloat16*>(out)[at] = __float2bfloat16_rn(EPI == 1 ? relu2(v) : v); + } +} + +// One 32-input group into acc: NT n8 row tiles, B fragments by `ldb(nt, s)` (uint2: b0, b1). +template +__device__ __forceinline__ void group(float (&acc)[NT][4], const uint4& w, uint32_t dp, int t, const uint2* tab, + LDB ldb) { + uint32_t a[2][4]; + float sc[2]; + F::decode(w, dp, t, tab, a, sc); + float p[NT][4]; +#pragma unroll + for (int nt = 0; nt < NT; ++nt) p[nt][0] = p[nt][1] = p[nt][2] = p[nt][3] = 0.f; +#pragma unroll + for (int s = 0; s < 2; ++s) +#pragma unroll + for (int nt = 0; nt < NT; ++nt) { + const uint2 b = ldb(nt, s); + mma(p[nt], a[s][0], a[s][1], a[s][2], a[s][3], b.x, b.y); + } +#pragma unroll + for (int nt = 0; nt < NT; ++nt) { + acc[nt][0] = fmaf(p[nt][0], sc[0], acc[nt][0]); + acc[nt][1] = fmaf(p[nt][1], sc[0], acc[nt][1]); + acc[nt][2] = fmaf(p[nt][2], sc[1], acc[nt][2]); + acc[nt][3] = fmaf(p[nt][3], sc[1], acc[nt][3]); + } +} + +// acc[nt] holds columns col0 + gq (0, 1) and col0 + gq + 8 (2, 3) of rows nt * 8 + 2t (0, 2) and + 1 (1, 3). +template +__device__ __forceinline__ void epilogue_rows(const float (&acc)[NT][4], void* out, int N, int col0, int lane, + const int* __restrict__ members, const Item& it) { + const int gq = lane >> 2, t = lane & 3; +#pragma unroll + for (int nt = 0; nt < NT; ++nt) +#pragma unroll + for (int j = 0; j < 2; ++j) { + const int r = nt * 8 + 2 * t + j; + if (r >= it.cnt) continue; + const size_t row = (size_t)pair_of(members, it, r) * N; + store(out, row + col0 + gq, acc[nt][j]); + store(out, row + col0 + gq + 8, acc[nt][2 + j]); + } +} + +constexpr int U = 2; // loads a warp has in flight, twice over + +// Decode rows: a CTA per TPC (item, 16-column tile) units, KW warps each over K's loads, summed in a fixed order. +template +__device__ __forceinline__ void gemv_body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint8_t* __restrict__ W, int KG, int NT16, + const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, int rows, + int skip) { + __shared__ uint2 tab[F::TABLE ? TABLE_WORDS : 1]; + __shared__ float red[TPC][KW > 1 ? KW - 1 : 1][8][32]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, gq = lane >> 2, t = lane & 3; + const int sub = warp / KW, kw = warp - sub * KW; + const int units = (items ? __ldg(counts) : 1) * NT16; + if (blockIdx.x * TPC >= units) return; + const int unit = blockIdx.x * TPC + sub, iu = unit / NT16, tile = unit - iu * NT16; + Item it = {0, 0, 0}; + if (unit < units) it = items ? Item{__ldg(items + 3 * iu), __ldg(items + 3 * iu + 1), __ldg(items + 3 * iu + 2)} + : Item{0, 0, rows}; + const bool live = unit < units && it.e != skip; + const uint8_t* wt = W + ((size_t)it.e * NT16 + tile) * (size_t)(KG / F::G) * F::BYTES; + const int KQ = KG / F::Q; + uint4 w[U]; + uint32_t dp[U]; + auto fetch = [&](int q0) { +#pragma unroll + for (int u = 0; u < U; ++u) + if (q0 + u * KW < KQ) { + w[u] = ld_nc(F::uptr(wt, q0 + u * KW, lane)); + dp[u] = ld_nc32(F::dptr(wt, (q0 + u * KW) * F::Q, lane)); + } + }; + if (live) fetch(kw); // the weights are on their way while the tables load + load_tables(tab); + const bool v0 = gq < it.cnt, v1 = gq + 8 < it.cnt, two = it.cnt > 8; + const int p0 = v0 ? pair_of(members, it, gq) : 0, p1 = v1 ? pair_of(members, it, gq + 8) : 0; + const __nv_bfloat16* x0 = X + (size_t)(slots ? p0 / slots : p0) * x_stride + 8 * t; + const __nv_bfloat16* x1 = X + (size_t)(slots ? p1 / slots : p1) * x_stride + 8 * t; + float acc[2][4]; +#pragma unroll + for (int nt = 0; nt < 2; ++nt) acc[nt][0] = acc[nt][1] = acc[nt][2] = acc[nt][3] = 0.f; + // a load's rows come in at once, Q groups' b fragments (steps 0 and 1) for row tiles 0 and 1 + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + auto xrows = [&](const __nv_bfloat16* x, bool v, int g) { + return v ? __ldg(reinterpret_cast(x + g * 32)) : zero; + }; + for (int q = live ? kw : KQ; q < KQ; q += U * KW) { + uint4 cw[U]; + uint32_t cd[U]; +#pragma unroll + for (int u = 0; u < U; ++u) { + cw[u] = w[u]; + cd[u] = dp[u]; + } + fetch(q + U * KW); +#pragma unroll + for (int u = 0; u < U; ++u) { + const int qu = q + u * KW; + if (qu >= KQ) break; + uint4 xr[F::Q][2]; +#pragma unroll + for (int j = 0; j < F::Q; ++j) { + xr[j][0] = xrows(x0, v0, qu * F::Q + j); + xr[j][1] = two ? xrows(x1, v1, qu * F::Q + j) : zero; + } +#pragma unroll + for (int j = 0; j < F::Q; ++j) { + uint4 wv = cw[u]; + if constexpr (F::Q > 1) { // group 4qu + j of columns gq, gq + 8 sits in lane (gq, j) + const int src = (lane & ~3) | j; + wv = make_uint4(__shfl_sync(~0u, wv.x, src), __shfl_sync(~0u, wv.y, src), __shfl_sync(~0u, wv.z, src), + __shfl_sync(~0u, wv.w, src)); + } + auto ldb = [&](int nt, int s) { + const uint4& v = xr[j][nt]; + return s ? make_uint2(v.z, v.w) : make_uint2(v.x, v.y); + }; + if (two) { + group(acc, wv, cd[u], t, tab, ldb); + } else { + float(&a1)[1][4] = *reinterpret_cast(&acc[0]); + group(a1, wv, cd[u], t, tab, ldb); + } + } + } + } + if constexpr (KW > 1) { + if (kw > 0) +#pragma unroll + for (int q = 0; q < 8; ++q) red[sub][kw - 1][q][lane] = acc[q >> 2][q & 3]; + __syncthreads(); + if (kw > 0) return; + for (int wi = 0; wi < KW - 1; ++wi) +#pragma unroll + for (int q = 0; q < 8; ++q) acc[q >> 2][q & 3] += red[sub][wi][q][lane]; + } + if (!live) return; + epilogue_rows(acc, out, N, tile * 16, lane, members, it); +} + +// Prompt rows: a CTA per (item of up to 64 rows, WN 16-column tiles), rows staged 128 inputs a stage. +constexpr int BM = 64, SK = 128, XC = SK / 8, STAGES = 3; // XC: 16-byte chunks a staged row; 48 KiB of stages + +__device__ __forceinline__ int chunk(int r, int c) { return r * XC + (c ^ (r & 7)); } + +template +__device__ __forceinline__ void prompt_tiles(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint8_t* __restrict__ W, int KG, int NT16, + const int* __restrict__ members, void* __restrict__ out, int N, + const Item& it, int ct, const uint2* tab) { + constexpr int THREADS = WN * 32, XPT = BM * XC / THREADS; + extern __shared__ uint4 sm[]; + const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5, gq = lane >> 2, t = lane & 3; + const int tile = ct * WN + warp; + const bool active = tile < NT16; + const __nv_bfloat16* xsrc[XPT]; + int xdst[XPT]; +#pragma unroll + for (int i = 0; i < XPT; ++i) { + const int q = tid + i * THREADS, r = q / XC, c = q - r * XC; + xdst[i] = chunk(r, c); + xsrc[i] = nullptr; + if (r < it.cnt) { + const int p = pair_of(members, it, r); + xsrc[i] = X + (size_t)(slots ? p / slots : p) * x_stride + 8 * c; + } + } + auto stage = [&](int s, int k) { + uint4* xs = sm + s * BM * XC; +#pragma unroll + for (int i = 0; i < XPT; ++i) + if (xsrc[i]) cp16(xs + xdst[i], xsrc[i] + (size_t)k * SK); + }; + const uint8_t* wt = W + ((size_t)it.e * NT16 + (active ? tile : 0)) * (size_t)(KG / F::G) * F::BYTES; + constexpr int GS = SK / 32; // groups a stage + float acc[NT][4]; +#pragma unroll + for (int nt = 0; nt < NT; ++nt) acc[nt][0] = acc[nt][1] = acc[nt][2] = acc[nt][3] = 0.f; + const int steps = KG / GS; +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) { + if (s < steps) stage(s, s); + cp_commit(); + } + uint4 w[GS]; + uint32_t dp[GS]; + if (active) +#pragma unroll + for (int gi = 0; gi < GS; ++gi) { + w[gi] = ld_nc(F::wptr(wt, gi, lane)); + dp[gi] = ld_nc32(F::dptr(wt, gi, lane)); + } + for (int k = 0; k < steps; ++k) { + cp_wait(); + __syncthreads(); + if (k + STAGES - 1 < steps) stage((k + STAGES - 1) % STAGES, k + STAGES - 1); + cp_commit(); + if (!active) continue; + uint4 wn[GS]; + uint32_t dn[GS]; + if (k + 1 < steps) +#pragma unroll + for (int gi = 0; gi < GS; ++gi) { + wn[gi] = ld_nc(F::wptr(wt, (k + 1) * GS + gi, lane)); + dn[gi] = ld_nc32(F::dptr(wt, (k + 1) * GS + gi, lane)); + } + const uint4* xs = sm + (k % STAGES) * BM * XC; +#pragma unroll + for (int gi = 0; gi < GS; ++gi) + group(acc, w[gi], dp[gi], t, tab, [&](int nt, int s) { + const int r = nt * 8 + gq; + return reinterpret_cast(xs + chunk(r, gi * 4 + t))[s]; + }); +#pragma unroll + for (int gi = 0; gi < GS; ++gi) { + w[gi] = wn[gi]; + dp[gi] = dn[gi]; + } + } + cp_wait<0>(); + if (active) epilogue_rows(acc, out, N, tile * 16, lane, members, it); +} + +template +__device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X, int x_stride, int slots, + const uint8_t* __restrict__ W, int KG, int NT16, + const int* __restrict__ items, const int* __restrict__ counts, + const int* __restrict__ members, void* __restrict__ out, int N, int rows, + int skip) { + __shared__ uint2 tab[F::TABLE ? TABLE_WORDS : 1]; + const int nct = (NT16 + WN - 1) / WN, iu = blockIdx.x / nct, ct = blockIdx.x - iu * nct; + Item it; + if (items) { + if (iu >= __ldg(counts)) return; + it = {__ldg(items + 3 * iu), __ldg(items + 3 * iu + 1), __ldg(items + 3 * iu + 2)}; + if (it.e == skip) return; + } else { + it = {0, iu * BM, min(BM, rows - iu * BM)}; + } + load_tables(tab); + const int nt = min(8, (it.cnt + 7) >> 3); +#define TF_TILES(NT) \ + prompt_tiles(X, x_stride, slots, W, KG, NT16, members, out, N, it, ct, tab) + switch (nt) { + case 8: TF_TILES(8); break; + case 7: TF_TILES(7); break; + case 6: TF_TILES(6); break; + case 5: TF_TILES(5); break; + case 4: TF_TILES(4); break; + case 3: TF_TILES(3); break; + case 2: TF_TILES(2); break; + default: TF_TILES(1); + } +#undef TF_TILES +} + +} // namespace tf_ggml + +// Entries: format x (decode KW 4, decode TPC 4, prompt) x epilogue (0 fp32, 1 relu^2 bf16, 3 bf16); dense: no items. +#define TF_GGML_GEMV(NAME, F, EPI, KW, TPC) \ + extern "C" __global__ void __launch_bounds__(KW * TPC * 32) NAME( \ + const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint8_t* __restrict__ W, int KG, int NT16, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int rows, int skip) { \ + tf_ggml::gemv_body(X, x_stride, slots, W, KG, NT16, items, counts, members, out, N, \ + rows, skip); \ + } +#define TF_GGML_ENTRIES(FMT, F, EPI, SUFFIX) \ + TF_GGML_GEMV(tf_ggml_##FMT##_gemv_##SUFFIX, F, EPI, 4, 1) \ + TF_GGML_GEMV(tf_ggml_##FMT##_gemv4_##SUFFIX, F, EPI, 1, 4) \ + extern "C" __global__ void __launch_bounds__(256) tf_ggml_##FMT##_prompt_##SUFFIX( \ + const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint8_t* __restrict__ W, int KG, int NT16, \ + const int* __restrict__ items, const int* __restrict__ counts, const int* __restrict__ members, \ + void* __restrict__ out, int N, int rows, int skip) { \ + tf_ggml::prompt_body(X, x_stride, slots, W, KG, NT16, items, counts, members, out, N, rows, \ + skip); \ + } +TF_GGML_ENTRIES(iq2, Iq2, 1, relu2) +TF_GGML_ENTRIES(iq2, Iq2, 0, f32) +TF_GGML_ENTRIES(iq2, Iq2, 3, bf16) +TF_GGML_ENTRIES(q8, Q8, 1, relu2) +TF_GGML_ENTRIES(q8, Q8, 0, f32) +TF_GGML_ENTRIES(q8, Q8, 3, bf16) +#undef TF_GGML_ENTRIES +#undef TF_GGML_GEMV + +// Load repack, a byte a thread: `count` matrices [n][k / KB][B] of checkpoint blocks into the 16-column tiles. +extern "C" __global__ void tf_ggml_pack(const uint8_t* __restrict__ src, uint8_t* __restrict__ dst, int iq2, int n, + int kb, long long total) { + const int B = iq2 ? 66 : 34, TB = iq2 ? 1056 : 544; + const long long per = (long long)n * kb * B; + for (long long o = blockIdx.x * (long long)blockDim.x + threadIdx.x; o < total; o += (long long)gridDim.x * blockDim.x) { + const long long m = o / per, rem = o - m * per; + const long long tb = rem / TB; + const int r = (int)(rem - tb * TB), tile = (int)(tb / kb), b = (int)(tb - (long long)tile * kb); + int col, at; + if (r >= TB - 32) { + const int q = r - (TB - 32); + col = tile * 16 + (q >> 2) + 8 * ((q >> 1) & 1); + at = q & 1; + } else if (iq2) { + col = tile * 16 + ((r >> 4) & 7) + 8 * ((r >> 3) & 1); + at = 2 + (r >> 7) * 8 + (r & 7); + } else { + const int ln = r >> 4, q = (r >> 2) & 3; + col = tile * 16 + (ln >> 2) + 8 * (q & 1); + at = 2 + 8 * (ln & 3) + 4 * (q >> 1) + (r & 3); + } + dst[o] = src[m * per + ((long long)col * kb + b) * B + at]; + } +} diff --git a/zig/kernels/cuda/ggml_iq2xxs_tables.cuh b/zig/kernels/cuda/ggml_iq2xxs_tables.cuh new file mode 100644 index 0000000000..f3f3b4e2b1 --- /dev/null +++ b/zig/kernels/cuda/ggml_iq2xxs_tables.cuh @@ -0,0 +1,82 @@ +// IQ2_XXS tables from ggml-common.h, MIT License, Copyright (c) 2023-2026 The ggml authors; also read by bolt_ggml.py. +#pragma once +#include + +static __device__ __constant__ const uint8_t tf_ksigns_iq2xs[128] = { + 0, 129, 130, 3, 132, 5, 6, 135, 136, 9, 10, 139, 12, 141, 142, 15, + 144, 17, 18, 147, 20, 149, 150, 23, 24, 153, 154, 27, 156, 29, 30, 159, + 160, 33, 34, 163, 36, 165, 166, 39, 40, 169, 170, 43, 172, 45, 46, 175, + 48, 177, 178, 51, 180, 53, 54, 183, 184, 57, 58, 187, 60, 189, 190, 63, + 192, 65, 66, 195, 68, 197, 198, 71, 72, 201, 202, 75, 204, 77, 78, 207, + 80, 209, 210, 83, 212, 85, 86, 215, 216, 89, 90, 219, 92, 221, 222, 95, + 96, 225, 226, 99, 228, 101, 102, 231, 232, 105, 106, 235, 108, 237, 238, 111, + 240, 113, 114, 243, 116, 245, 246, 119, 120, 249, 250, 123, 252, 125, 126, 255, +}; + +// In global memory: kernels copy it to shared memory with coalesced loads. +static __device__ const uint64_t tf_iq2xxs_grid[256] = { + 0x0808080808080808, 0x080808080808082b, 0x0808080808081919, 0x0808080808082b08, + 0x0808080808082b2b, 0x0808080808190819, 0x0808080808191908, 0x08080808082b0808, + 0x08080808082b082b, 0x08080808082b2b08, 0x08080808082b2b2b, 0x0808080819080819, + 0x0808080819081908, 0x0808080819190808, 0x0808080819192b08, 0x08080808192b0819, + 0x08080808192b1908, 0x080808082b080808, 0x080808082b08082b, 0x080808082b082b2b, + 0x080808082b2b082b, 0x0808081908080819, 0x0808081908081908, 0x0808081908190808, + 0x0808081908191919, 0x0808081919080808, 0x080808192b081908, 0x080808192b192b08, + 0x0808082b08080808, 0x0808082b0808082b, 0x0808082b082b082b, 0x0808082b2b08082b, + 0x0808190808080819, 0x0808190808081908, 0x0808190808190808, 0x08081908082b0819, + 0x08081908082b1908, 0x0808190819080808, 0x080819081908082b, 0x0808190819082b08, + 0x08081908192b0808, 0x080819082b080819, 0x080819082b081908, 0x080819082b190808, + 0x080819082b2b1908, 0x0808191908080808, 0x080819190808082b, 0x0808191908082b08, + 0x08081919082b0808, 0x080819191908192b, 0x08081919192b2b19, 0x080819192b080808, + 0x080819192b190819, 0x0808192b08082b19, 0x0808192b08190808, 0x0808192b19080808, + 0x0808192b2b081908, 0x0808192b2b2b1908, 0x08082b0808080808, 0x08082b0808081919, + 0x08082b0808082b08, 0x08082b0808191908, 0x08082b08082b2b08, 0x08082b0819080819, + 0x08082b0819081908, 0x08082b0819190808, 0x08082b081919082b, 0x08082b082b082b08, + 0x08082b1908081908, 0x08082b1919080808, 0x08082b2b0808082b, 0x08082b2b08191908, + 0x0819080808080819, 0x0819080808081908, 0x0819080808190808, 0x08190808082b0819, + 0x0819080819080808, 0x08190808192b0808, 0x081908082b081908, 0x081908082b190808, + 0x081908082b191919, 0x0819081908080808, 0x0819081908082b08, 0x08190819082b0808, + 0x0819081919190808, 0x0819081919192b2b, 0x081908192b080808, 0x0819082b082b1908, + 0x0819082b19081919, 0x0819190808080808, 0x0819190808082b08, 0x08191908082b0808, + 0x08191908082b1919, 0x0819190819082b19, 0x081919082b080808, 0x0819191908192b08, + 0x08191919192b082b, 0x0819192b08080808, 0x0819192b0819192b, 0x08192b0808080819, + 0x08192b0808081908, 0x08192b0808190808, 0x08192b0819080808, 0x08192b082b080819, + 0x08192b1908080808, 0x08192b1908081919, 0x08192b192b2b0808, 0x08192b2b19190819, + 0x082b080808080808, 0x082b08080808082b, 0x082b080808082b2b, 0x082b080819081908, + 0x082b0808192b0819, 0x082b08082b080808, 0x082b08082b08082b, 0x082b0819082b2b19, + 0x082b081919082b08, 0x082b082b08080808, 0x082b082b0808082b, 0x082b190808080819, + 0x082b190808081908, 0x082b190808190808, 0x082b190819080808, 0x082b19081919192b, + 0x082b191908080808, 0x082b191919080819, 0x082b1919192b1908, 0x082b192b2b190808, + 0x082b2b0808082b08, 0x082b2b08082b0808, 0x082b2b082b191908, 0x082b2b2b19081908, + 0x1908080808080819, 0x1908080808081908, 0x1908080808190808, 0x1908080808192b08, + 0x19080808082b0819, 0x19080808082b1908, 0x1908080819080808, 0x1908080819082b08, + 0x190808081919192b, 0x19080808192b0808, 0x190808082b080819, 0x190808082b081908, + 0x190808082b190808, 0x1908081908080808, 0x19080819082b0808, 0x19080819192b0819, + 0x190808192b080808, 0x190808192b081919, 0x1908082b08080819, 0x1908082b08190808, + 0x1908082b19082b08, 0x1908082b1919192b, 0x1908082b192b2b08, 0x1908190808080808, + 0x1908190808082b08, 0x19081908082b0808, 0x190819082b080808, 0x190819082b192b19, + 0x190819190819082b, 0x19081919082b1908, 0x1908192b08080808, 0x19082b0808080819, + 0x19082b0808081908, 0x19082b0808190808, 0x19082b0819080808, 0x19082b0819081919, + 0x19082b1908080808, 0x19082b1919192b08, 0x19082b19192b0819, 0x19082b192b08082b, + 0x19082b2b19081919, 0x19082b2b2b190808, 0x1919080808080808, 0x1919080808082b08, + 0x1919080808190819, 0x1919080808192b19, 0x19190808082b0808, 0x191908082b080808, + 0x191908082b082b08, 0x1919081908081908, 0x191908191908082b, 0x191908192b2b1908, + 0x1919082b2b190819, 0x191919082b190808, 0x191919082b19082b, 0x1919191908082b2b, + 0x1919192b08080819, 0x1919192b19191908, 0x19192b0808080808, 0x19192b0808190819, + 0x19192b0808192b19, 0x19192b08192b1908, 0x19192b1919080808, 0x19192b2b08082b08, + 0x192b080808081908, 0x192b080808190808, 0x192b080819080808, 0x192b0808192b2b08, + 0x192b081908080808, 0x192b081919191919, 0x192b082b08192b08, 0x192b082b192b0808, + 0x192b190808080808, 0x192b190808081919, 0x192b191908190808, 0x192b19190819082b, + 0x192b19192b081908, 0x192b2b081908082b, 0x2b08080808080808, 0x2b0808080808082b, + 0x2b08080808082b2b, 0x2b08080819080819, 0x2b0808082b08082b, 0x2b08081908081908, + 0x2b08081908192b08, 0x2b08081919080808, 0x2b08082b08190819, 0x2b08190808080819, + 0x2b08190808081908, 0x2b08190808190808, 0x2b08190808191919, 0x2b08190819080808, + 0x2b081908192b0808, 0x2b08191908080808, 0x2b0819191908192b, 0x2b0819192b191908, + 0x2b08192b08082b19, 0x2b08192b19080808, 0x2b08192b192b0808, 0x2b082b080808082b, + 0x2b082b1908081908, 0x2b082b2b08190819, 0x2b19080808081908, 0x2b19080808190808, + 0x2b190808082b1908, 0x2b19080819080808, 0x2b1908082b2b0819, 0x2b1908190819192b, + 0x2b1908192b080808, 0x2b19082b19081919, 0x2b19190808080808, 0x2b191908082b082b, + 0x2b19190819081908, 0x2b19191919190819, 0x2b192b082b080819, 0x2b192b19082b0808, + 0x2b2b08080808082b, 0x2b2b080819190808, 0x2b2b08082b081919, 0x2b2b081908082b19, + 0x2b2b082b08080808, 0x2b2b190808192b08, 0x2b2b2b0819190808, 0x2b2b2b1908081908, +}; diff --git a/zig/kernels/cuda/nemotron_ops.cu b/zig/kernels/cuda/nemotron_ops.cu index b8bd0603de..b55bbc712d 100644 --- a/zig/kernels/cuda/nemotron_ops.cu +++ b/zig/kernels/cuda/nemotron_ops.cu @@ -291,3 +291,31 @@ extern "C" __global__ void tf_add_row(void* __restrict__ y, int y_f32, const flo *p = __float2bfloat16(__bfloat162float(*p) + x[i]); } } + +// LatentMoE (Bolt): Z [R, L] bf16, the fp32 sum of Y [R * K, L] (slot order) times WT [R * K]; one block a row. +extern "C" __global__ void tf_latent_sum(const float* __restrict__ Y, const float* __restrict__ WT, __nv_bfloat16* __restrict__ Z, + int K, int L) { + const int64_t r = blockIdx.x; + for (int i = threadIdx.x; i < L; i += blockDim.x) { + float acc = 0.0f; + for (int k = 0; k < K; ++k) acc = __fmaf_rn(Y[(r * K + k) * L + i], WT[r * K + k], acc); + Z[r * L + i] = __float2bfloat16(acc); + } +} + +// out = a + b in bf16 (b bf16, or fp32 when b_f32): the latent experts' projection plus the shared expert, n values. +extern "C" __global__ void tf_add_bf16(__nv_bfloat16* __restrict__ out, const __nv_bfloat16* __restrict__ a, const void* __restrict__ b, + int b_f32, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + const float y = b_f32 ? static_cast(b)[i] : __bfloat162float(static_cast(b)[i]); + out[i] = __float2bfloat16(__bfloat162float(a[i]) + y); +} + +// Image prompt chunk: row r of H [rows, D] takes F's row SRC[r] for its token embedding (-1: text); a block a row. +extern "C" __global__ void tf_splice_rows(__nv_bfloat16* __restrict__ H, const __nv_bfloat16* __restrict__ F, const int* __restrict__ SRC, int D) { + const int64_t r = blockIdx.x; + const int s = SRC[r]; + if (s < 0) return; + for (int i = threadIdx.x; i < D; i += blockDim.x) H[r * D + i] = F[static_cast(s) * D + i]; +} diff --git a/zig/kernels/cuda/nemotron_route.cu b/zig/kernels/cuda/nemotron_route.cu index 23c6b78f35..4538440a5e 100644 --- a/zig/kernels/cuda/nemotron_route.cu +++ b/zig/kernels/cuda/nemotron_route.cu @@ -38,7 +38,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_nemo_router(const __nv_bflo namespace { -constexpr int MAXC = 8; // experts a lane holds: up to 256 experts +constexpr int MAXC = 16; // experts a lane holds: up to 512 experts (Bolt's widest layer has 496) // The larger of two (value, index) picks: a higher value, or the lower index on a tie (values are never NaN). __device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return v > bv || (v == bv && i < bi); } diff --git a/zig/kernels/cuda/radio_vision.cu b/zig/kernels/cuda/radio_vision.cu new file mode 100644 index 0000000000..f930c54eb6 --- /dev/null +++ b/zig/kernels/cuda/radio_vision.cu @@ -0,0 +1,269 @@ +// Bolt's C-RADIO ViT-H/16 tower and projector around cuBLASLt bf16 GEMMs; residual, norms, softmax, epilogues fp32. +#include +#include + +__device__ __forceinline__ float bf(__nv_bfloat16 v) { return __bfloat162float(v); } + +__device__ __forceinline__ uint32_t pack2(float lo, float hi) { + __nv_bfloat162 h = __floats2bfloat162_rn(lo, hi); + return *reinterpret_cast(&h); +} + +// fp32 block sum over blockDim.x threads (a multiple of 32, at most 1024); every thread gets the total. +__device__ __forceinline__ float block_sum(float v, float* part) { + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5, warps = blockDim.x >> 5; + for (int o = 16; o; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o); + __syncthreads(); + if (lane == 0) part[warp] = v; + __syncthreads(); + float t = 0.0f; + for (int i = 0; i < warps; ++i) t += part[i]; + return t; +} + +// x [prefix + n, D] fp32: cls/register rows, then acc (patch Linear) plus CPE rows bilinear from the 128x128 table. +extern "C" __global__ void __launch_bounds__(256) tf_rv_embed(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ pos, + const int* __restrict__ idx, + const float* __restrict__ wts, + const __nv_bfloat16* __restrict__ cls, + float* __restrict__ x, int prefix, int D) { + const long long row = blockIdx.x; + if (row < prefix) { + for (int i = threadIdx.x; i < D; i += blockDim.x) x[row * D + i] = bf(cls[row * D + i]); + return; + } + const long long p = row - prefix; + const int* ri = idx + p * 4; + const float* rw = wts + p * 4; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float a = bf(pos[(long long)ri[0] * D + i]) * rw[0] + bf(pos[(long long)ri[1] * D + i]) * rw[1]; + const float b = bf(pos[(long long)ri[2] * D + i]) * rw[0] + bf(pos[(long long)ri[3] * D + i]) * rw[1]; + x[row * D + i] = acc[p * D + i] + (a * rw[2] + b * rw[3]); + } +} + +// LayerNorm of fp32 rows (mean, then the biased variance about it), bf16 out for the next GEMM; a block a row. +extern "C" __global__ void __launch_bounds__(256) tf_rv_layernorm(const float* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ b, + __nv_bfloat16* __restrict__ y, int D, float eps) { + __shared__ float part[32]; + const long long row = blockIdx.x; + const float* xr = x + row * D; + float s = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) s += xr[i]; + const float mean = block_sum(s, part) / D; + float v = 0.0f; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float d = xr[i] - mean; + v += d * d; + } + const float r = rsqrtf(block_sum(v, part) / D + eps); + for (int i = threadIdx.x; i < D; i += blockDim.x) + y[row * D + i] = __float2bfloat16_rn((xr[i] - mean) * r * bf(w[i]) + bf(b[i])); +} + +// acc [rows, 3D] + bias split into q, k, v [heads][rows][hd] bf16 (the reference's view(T, 3, nh, hd)). +extern "C" __global__ void __launch_bounds__(256) tf_rv_qkv(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ q, __nv_bfloat16* __restrict__ k, + __nv_bfloat16* __restrict__ v, int rows, int heads, int hd) { + const long long row = blockIdx.x; + const int D = heads * hd; + const float* a = acc + row * 3 * D; + for (int c = threadIdx.x; c < 3 * D; c += blockDim.x) { + const int s = c / D, h = (c % D) / hd, j = c % hd; + const __nv_bfloat16 val = __float2bfloat16_rn(a[c] + bf(bias[c])); + __nv_bfloat16* dst = s == 0 ? q : (s == 1 ? k : v); + dst[((long long)h * rows + row) * hd + j] = val; + } +} + +__device__ __forceinline__ void mma16816(float* c, const uint32_t* a, uint32_t b0, uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, " + "{%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +} + +// Bidirectional flash attention, heads of 80: a warp 16 queries, 64-key tiles, bf16 mma, fp32 softmax; bf16 out. +#define RV_HD 80 +#define RV_BK 64 +extern "C" __global__ void __launch_bounds__(128) tf_rv_attention(const __nv_bfloat16* __restrict__ q, + const __nv_bfloat16* __restrict__ k, + const __nv_bfloat16* __restrict__ v, + __nv_bfloat16* __restrict__ out, int rows, int heads, + float scale_log2) { + __shared__ __align__(16) __nv_bfloat16 ks[RV_BK][RV_HD + 8]; + __shared__ __align__(16) __nv_bfloat16 vt[RV_HD][RV_BK + 8]; + const int head = blockIdx.y, warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int g = lane >> 2, t = lane & 3; + const long long base = (long long)head * rows; + const int q0 = blockIdx.x * 64 + warp * 16; + uint32_t qa[RV_HD / 16][4]; + { + const uint32_t* p0 = reinterpret_cast(q + (base + min(q0 + g, rows - 1)) * RV_HD); + const uint32_t* p1 = reinterpret_cast(q + (base + min(q0 + g + 8, rows - 1)) * RV_HD); +#pragma unroll + for (int kc = 0; kc < RV_HD / 16; ++kc) { + qa[kc][0] = p0[kc * 8 + t]; + qa[kc][1] = p1[kc * 8 + t]; + qa[kc][2] = p0[kc * 8 + 4 + t]; + qa[kc][3] = p1[kc * 8 + 4 + t]; + } + } + float o[RV_HD / 8][4]; +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) o[d][0] = o[d][1] = o[d][2] = o[d][3] = 0.0f; + float m0 = -INFINITY, m1 = -INFINITY, l0 = 0.0f, l1 = 0.0f; + for (int k0 = 0; k0 < rows; k0 += RV_BK) { + __syncthreads(); + for (int i = threadIdx.x; i < RV_BK * RV_HD / 2; i += blockDim.x) { + const int r = i / (RV_HD / 2), c = i % (RV_HD / 2); + uint32_t kv = 0, vv = 0; + if (k0 + r < rows) { + kv = reinterpret_cast(k + (base + k0 + r) * RV_HD)[c]; + vv = reinterpret_cast(v + (base + k0 + r) * RV_HD)[c]; + } + *reinterpret_cast(&ks[r][2 * c]) = kv; + const __nv_bfloat162 vh = *reinterpret_cast<__nv_bfloat162*>(&vv); + vt[2 * c][r] = vh.x; + vt[2 * c + 1][r] = vh.y; + } + __syncthreads(); + float s[RV_BK / 8][4]; +#pragma unroll + for (int nt = 0; nt < RV_BK / 8; ++nt) { + s[nt][0] = s[nt][1] = s[nt][2] = s[nt][3] = 0.0f; + const uint32_t* kr = reinterpret_cast(&ks[nt * 8 + g][0]); +#pragma unroll + for (int kc = 0; kc < RV_HD / 16; ++kc) mma16816(s[nt], qa[kc], kr[kc * 8 + t], kr[kc * 8 + 4 + t]); + } + float mx0 = -INFINITY, mx1 = -INFINITY; +#pragma unroll + for (int nt = 0; nt < RV_BK / 8; ++nt) { +#pragma unroll + for (int e = 0; e < 4; ++e) { + const int key = k0 + nt * 8 + 2 * t + (e & 1); + s[nt][e] = key < rows ? s[nt][e] * scale_log2 : -INFINITY; + } + mx0 = fmaxf(mx0, fmaxf(s[nt][0], s[nt][1])); + mx1 = fmaxf(mx1, fmaxf(s[nt][2], s[nt][3])); + } + for (int off = 1; off < 4; off <<= 1) { + mx0 = fmaxf(mx0, __shfl_xor_sync(0xffffffffu, mx0, off)); + mx1 = fmaxf(mx1, __shfl_xor_sync(0xffffffffu, mx1, off)); + } + const float n0 = fmaxf(m0, mx0), n1 = fmaxf(m1, mx1); + const float c0 = exp2f(m0 - n0), c1 = exp2f(m1 - n1); + m0 = n0; + m1 = n1; + float sum0 = 0.0f, sum1 = 0.0f; +#pragma unroll + for (int nt = 0; nt < RV_BK / 8; ++nt) { + s[nt][0] = exp2f(s[nt][0] - n0); + s[nt][1] = exp2f(s[nt][1] - n0); + s[nt][2] = exp2f(s[nt][2] - n1); + s[nt][3] = exp2f(s[nt][3] - n1); + sum0 += s[nt][0] + s[nt][1]; + sum1 += s[nt][2] + s[nt][3]; + } + l0 = l0 * c0 + sum0; + l1 = l1 * c1 + sum1; +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) { + o[d][0] *= c0; + o[d][1] *= c0; + o[d][2] *= c1; + o[d][3] *= c1; + } +#pragma unroll + for (int kc = 0; kc < RV_BK / 16; ++kc) { + uint32_t pa[4]; + pa[0] = pack2(s[2 * kc][0], s[2 * kc][1]); + pa[1] = pack2(s[2 * kc][2], s[2 * kc][3]); + pa[2] = pack2(s[2 * kc + 1][0], s[2 * kc + 1][1]); + pa[3] = pack2(s[2 * kc + 1][2], s[2 * kc + 1][3]); +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) { + const uint32_t* vr = reinterpret_cast(&vt[d * 8 + g][0]); + mma16816(o[d], pa, vr[kc * 8 + t], vr[kc * 8 + 4 + t]); + } + } + } + for (int off = 1; off < 4; off <<= 1) { + l0 += __shfl_xor_sync(0xffffffffu, l0, off); + l1 += __shfl_xor_sync(0xffffffffu, l1, off); + } + const float i0 = 1.0f / l0, i1 = 1.0f / l1; + const int D = heads * RV_HD; + const int r0 = q0 + g, r1 = q0 + g + 8; +#pragma unroll + for (int d = 0; d < RV_HD / 8; ++d) { + const int col = head * RV_HD + d * 8 + 2 * t; + if (r0 < rows) *reinterpret_cast(out + (long long)r0 * D + col) = pack2(o[d][0] * i0, o[d][1] * i0); + if (r1 < rows) *reinterpret_cast(out + (long long)r1 * D + col) = pack2(o[d][2] * i1, o[d][3] * i1); + } +} + +// y = gelu(acc + b), the exact erf form (nn.GELU), fp32, bf16 out. +extern "C" __global__ void __launch_bounds__(256) tf_rv_bias_gelu(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ y, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float x = acc[row * D + i] + bf(bias[i]); + y[row * D + i] = __float2bfloat16_rn(0.5f * x * (1.0f + erff(x * 0.7071067811865476f))); + } +} + +// x += acc + b, the residual in fp32. +extern "C" __global__ void __launch_bounds__(256) tf_rv_bias_residual(const float* __restrict__ acc, + const __nv_bfloat16* __restrict__ bias, + float* __restrict__ x, int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) x[row * D + i] += acc[row * D + i] + bf(bias[i]); +} + +// Pixel shuffle v2 (2x2 patches to a 4D row) of the rows past `prefix`, then fp32 RMSNorm(4D) by w; a block a row. +extern "C" __global__ void __launch_bounds__(256) tf_rv_shuffle_rms(const float* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + __nv_bfloat16* __restrict__ y, int prefix, int gw, + int D, float eps) { + __shared__ float part[32]; + const int r = blockIdx.x; + const int by = r / (gw / 2), bx = r % (gw / 2); + const int W = 4 * D; + float s = 0.0f; + for (int ch = threadIdx.x; ch < W; ch += blockDim.x) { + const int sub = ch / D, c = ch % D; + const long long tok = prefix + (long long)(2 * by + sub / 2) * gw + 2 * bx + sub % 2; + const float v = x[tok * D + c]; + s += v * v; + } + const float rs = rsqrtf(block_sum(s, part) / W + eps); + for (int ch = threadIdx.x; ch < W; ch += blockDim.x) { + const int sub = ch / D, c = ch % D; + const long long tok = prefix + (long long)(2 * by + sub / 2) * gw + 2 * bx + sub % 2; + y[(long long)r * W + ch] = __float2bfloat16_rn(x[tok * D + c] * rs * bf(w[ch])); + } +} + +// y = relu(acc)^2, bf16 out. +extern "C" __global__ void __launch_bounds__(256) tf_rv_relu2(const float* __restrict__ acc, __nv_bfloat16* __restrict__ y, + int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float a = fmaxf(acc[row * D + i], 0.0f); + y[row * D + i] = __float2bfloat16_rn(a * a); + } +} + +// y = bf16(acc): the projector's last Linear. +extern "C" __global__ void __launch_bounds__(256) tf_rv_cast(const float* __restrict__ acc, __nv_bfloat16* __restrict__ y, + int D) { + const long long row = blockIdx.x; + for (int i = threadIdx.x; i < D; i += blockDim.x) y[row * D + i] = __float2bfloat16_rn(acc[row * D + i]); +} diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 64347ae50c..975db3e9a9 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -10,6 +10,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 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" }, diff --git a/zig/src/core/checkpoint.zig b/zig/src/core/checkpoint.zig index 1c2110a685..47ff2b5011 100644 --- a/zig/src/core/checkpoint.zig +++ b/zig/src/core/checkpoint.zig @@ -59,6 +59,13 @@ pub const Checkpoint = struct { io: Io, files: std.ArrayList(st.File) = .empty, used: std.StringHashMapUnmanaged(void) = .empty, + /// Prepended to the names get, has and skip take ("language_model." for Bolt's LM beside its vision tower). + scope: []const u8 = "", + + fn scoped(self: *const Checkpoint, buf: *[320]u8, name: []const u8) ![]const u8 { + if (self.scope.len == 0) return name; + return std.fmt.bufPrint(buf, "{s}{s}", .{ self.scope, name }); + } /// Every shard of `dir`, mapped. pub fn openModel(gpa: std.mem.Allocator, io: Io, dir: []const u8) !Checkpoint { @@ -90,7 +97,9 @@ pub const Checkpoint = struct { } /// The tensor named `name`, marked as used; a missing name is an error naming it. - pub fn get(self: *Checkpoint, name: []const u8) !Tensor { + pub fn get(self: *Checkpoint, short: []const u8) !Tensor { + var buf: [320]u8 = undefined; + const name = try self.scoped(&buf, short); for (self.files.items) |*f| if (f.get(name)) |t| { try self.used.put(self.gpa, f.names.getKey(name).?, {}); return t; @@ -110,13 +119,17 @@ pub const Checkpoint = struct { } /// Whether any file holds `name` (it is not marked as used). - pub fn has(self: *const Checkpoint, name: []const u8) bool { + pub fn has(self: *const Checkpoint, short: []const u8) bool { + var buf: [320]u8 = undefined; + const name = self.scoped(&buf, short) catch return false; for (self.files.items) |*f| if (f.get(name) != null) return true; return false; } /// Marks every name starting with `prefix` as used: a part the loader leaves out on purpose. - pub fn skip(self: *Checkpoint, prefix: []const u8) !void { + pub fn skip(self: *Checkpoint, short: []const u8) !void { + var buf: [320]u8 = undefined; + const prefix = try self.scoped(&buf, short); for (self.files.items) |*f| for (f.names.keys()) |k| { if (std.mem.startsWith(u8, k, prefix)) try self.used.put(self.gpa, k, {}); }; @@ -224,3 +237,26 @@ test "has finds a name without taking it; skip marks a prefix used" { _ = try ck.get("backbone.x"); try std.testing.expectEqual(@as(usize, 0), ck.unused()); } + +test "a scope reads a wrapper's part by its own names" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + for ([_][]const u8{ "language_model.backbone.x", "vision_model.y" }, 0..) |name, i| { + const image = try testImage(gpa, name, 4); + defer gpa.free(image); + var file: [32]u8 = undefined; + try tmp.dir.writeFile(io, .{ .sub_path = try std.fmt.bufPrint(&file, "model-{d}.safetensors", .{i}), .data = image }); + } + var path: [64]u8 = undefined; + var ck = try Checkpoint.openModel(gpa, io, try std.fmt.bufPrint(&path, ".zig-cache/tmp/{s}", .{tmp.sub_path})); + defer ck.close(); + ck.scope = "language_model."; + try std.testing.expect(ck.has("backbone.x") and !ck.has("vision_model.y")); + _ = try ck.get("backbone.x"); + try std.testing.expectEqual(@as(u32, 1), ck.used.count()); // vision_model.y is left + ck.scope = ""; + try ck.skip("vision_model."); + try std.testing.expectEqual(@as(usize, 0), ck.unused()); +} diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 903dc4d40b..1a86bbf3b4 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -1,5 +1,8 @@ //! The engine a server drives (``Engine``), and ``LaneHost``: the lane core served on one thread, rounds left to it. const std = @import("std"); +/// Image inputs (decode, resize, patches) for families with a vision tower; the server reaches them here. +pub const Image = lanes.Image; +pub const qwen_image = @import("qwen_image"); const lanes = @import("lanes"); const Allocator = std.mem.Allocator; @@ -88,6 +91,8 @@ pub const Request = struct { logprobs: ?u8 = null, /// Rows for each prompt token after the first, with ranks and this many best tokens, when ``Info.prompt_logprobs``. prompt_logprobs: ?u8 = null, + /// An image prompt's inputs in prompt order (Info.vision engines only; the prompt holds their placeholder rows). + images: []const lanes.Image = &.{}, }; pub const Reason = enum { stop, length, cancelled, failed }; @@ -162,6 +167,12 @@ pub const Info = struct { prompt_cache_plan: ?PromptCachePlan = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, + /// The engine encodes image inputs (Request.images): it was opened with ``vision`` and its family has a tower. + vision: bool = false, + /// Its tower takes GLM-5.3-Flash's preparation and placeholders (qwen_image.glm) instead of Qwen's. + vision_glm: bool = false, + /// Its tower takes Nemotron 3.5 Bolt's preparation (qwen_image.nemotron) and ... placeholders. + vision_nemotron: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. admit: ?Admit = null, }; @@ -225,10 +236,20 @@ pub const Open = struct { master_port: u16 = 29551, /// --policy: the engine's policy as `key=value,...`; empty: its defaults. policy: []const u8 = "", + /// --vision: load the family's vision tower and accept image inputs. + vision: bool = false, + /// --kv-dtype: bits of the attention cache's codes (8 or 4) for families that keep a quantized cache. + kv_bits: u8 = 8, }; /// An opened engine; ``close`` stops its thread and frees its backend. -pub const Opened = struct { engine: Engine, close: *const fn (ctx: *anyopaque) void, ctx: *anyopaque }; +pub const Opened = struct { + engine: Engine, + close: *const fn (ctx: *anyopaque) void, + ctx: *anyopaque, + /// Rank 1 of two: replay rank 0's engine calls until it stops (the server then serves no HTTP); null elsewhere. + follow: ?*const fn (ctx: *anyopaque) anyerror!void = null, +}; pub const Memory = struct { active: u64 = 0, cache: u64 = 0, peak: u64 = 0, pinned: u64 = 0 }; diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4394fdac05..73937e925c 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -85,6 +85,7 @@ pub const LaneHost = struct { const e = job.entry orelse return; job.entry = null; const store = job.host.cache orelse return; + defer store.unpin(e); // after resumed, which reads e if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); @@ -285,16 +286,16 @@ pub const LaneHost = struct { h.mutex.unlock(h.io); } - /// Prefills the next queued request into a free lane; false when none waits or no lane is free. - fn admitOne(h: *LaneHost) bool { + /// The next queued request into the admitted ones (null: none waits, no lane is free, or its caches do not fit). + pub fn takeJob(h: *LaneHost) ?*Job { h.lock(); if (h.queued.items.len == 0 or h.admitted.items.len >= @min(h.info_.lanes, h.memory_lanes orelse h.info_.lanes)) { h.unlock(); - return false; + return null; } const at = admit.pick(h) orelse { // a shared cache pool's head waits for room h.unlock(); - return false; + return null; }; const job = h.queued.orderedRemove(at); job.seq = h.admissions; @@ -302,20 +303,30 @@ pub const LaneHost = struct { h.admitted.append(h.gpa, job) catch { h.unlock(); h.finish(job, .failed, "out of memory"); - return true; + return null; }; h.unlock(); + return job; + } + + /// The admitted request's stream (its kept state, drafter and grammar); false when it was dropped instead. + pub fn setup(h: *LaneHost, job: *Job) bool { const r = job.request; - if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran + if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) { // other work came first: it never ran + _ = h.cancel(job); + return false; + } var reuse: lanes.stream.Reuse = .{}; - // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { - job.entry = l.entry; + job.entry = if (l.entry) |e| store.pin(e) else null; // kept alive until its restore reports job.kept0 = store.counts.kept; job.marks = l.marks; reuse = .{ .saved = if (l.entry) |e| e.saved else null, .at = if (l.entry) |e| e.at else 0, .marks = l.marks, .hook = .{ .ptr = job, .at = Job.kept } }; } else |_| {}; - job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch return h.drop(job, "the drafter could not start"); + job.proposer = lanes.SuffixLookup.init(h.gpa, .{ .min_match = h.min_match }) catch { + _ = h.drop(job, "the drafter could not start"); + return false; + }; job.stream = lanes.Stream.init(h.gpa, .{ .id = "request", .prompt = r.prompt, @@ -339,18 +350,31 @@ pub const LaneHost = struct { .reuse = reuse, .grammar = if (r.structure) |st| h.constraint(st) catch |e| { job.proposer.deinit(); - return h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + { + _ = h.drop(job, if (e == error.OutOfMemory) "out of memory" else "the reply's grammar could not start (xgrammar)"); + return false; + } } else null, + .images = r.images, .decode_spans = h.spans(r), .decide = r.decide, .untaught = r.untaught, }) catch { job.proposer.deinit(); - return h.drop(job, "out of memory"); + { + _ = h.drop(job, "out of memory"); + return false; + } }; job.started = true; - const began = std.Io.Clock.awake.now(h.io).toNanoseconds(); - job.began = began; + job.began = std.Io.Clock.awake.now(h.io).toNanoseconds(); + return true; + } + + /// An admitted request into the engine: a warm pass, the lone driver, the fill queue, or its own prompt pass. + pub fn launch(h: *LaneHost, job: *Job) bool { + const r = job.request; + const began = job.began; if (r.max_tokens == 0 and h.info_.warm_turns) return h.warmPass(job, began); if (lone.fits(h, job)) return lone.run(h, job, began, false); if (h.core.fills()) { @@ -444,12 +468,13 @@ pub const LaneHost = struct { } /// No memory for another lane yet: the job waits at the queue's head, and admission holds until a job ends. - fn requeue(h: *LaneHost, job: *Job) bool { + pub fn requeue(h: *LaneHost, job: *Job) bool { if (job.stream.grammar) |g| g.deinit(h.gpa); // the next admission compiles its own job.stream.deinit(h.gpa); job.proposer.deinit(); job.started = false; job.fill_began = false; + if (job.entry) |e| h.cache.?.unpin(e); job.entry = null; // looked up again at its next admission h.gpa.free(job.marks); job.marks = &.{}; @@ -475,7 +500,7 @@ pub const LaneHost = struct { const alone_words = "the engine has no memory for this request's caches with nothing else running: lower --context or --parallel"; /// A job cancelled in its prompt pass, its lane already released. - fn cancel(h: *LaneHost, job: *Job) bool { + pub fn cancel(h: *LaneHost, job: *Job) bool { h.remove(job); h.finish(job, .cancelled, ""); return true; @@ -494,7 +519,7 @@ pub const LaneHost = struct { fn run(h: *LaneHost) void { while (true) { h.takeCancels(); - while (h.admitOne()) {} + while (admit.many(h)) {} h.fillOne(); status.noteLive(h); h.lock(); diff --git a/zig/src/core/lane_host_admit.zig b/zig/src/core/lane_host_admit.zig index 016ab4398a..f394a3729e 100644 --- a/zig/src/core/lane_host_admit.zig +++ b/zig/src/core/lane_host_admit.zig @@ -2,6 +2,8 @@ const api = @import("engine_api.zig"); const LaneHost = @import("lane_host.zig").LaneHost; const Job = LaneHost.Job; +const lanes = @import("lanes"); +const lone = @import("lane_host_lone.zig"); /// The queued job to admit (its index), or null to hold; the caller holds the lock and a lane is free. pub fn pick(h: *LaneHost) ?usize { @@ -21,3 +23,51 @@ pub fn pick(h: *LaneHost) ?usize { fn fits(a: api.Admit, job: *const Job) bool { return a.fits(a.ctx, @intCast(job.request.prompt.len), job.request.max_tokens); } + +/// Every queued request a lane is free for, prompts in one pass where the backend takes several; else one at a time. +pub fn many(h: *LaneHost) bool { + if (!h.core.batchesPrompts() or h.core.fills()) return one(h); + h.lock(); + const waiting = h.queued.items.len; + const free = h.info_.lanes -| h.admitted.items.len; + h.unlock(); + if (waiting < 2 or free < 2) return one(h); + var jobs: [lanes.Engine.max_together]*Job = undefined; + var n: usize = 0; + while (n < jobs.len and n < free) { + const job = h.takeJob() orelse break; + if (!h.setup(job)) continue; + if ((job.request.max_tokens == 0 and h.info_.warm_turns) or lone.fits(h, job)) { + _ = h.launch(job); + continue; + } + jobs[n] = job; + n += 1; + } + if (n == 0) return true; + if (n == 1) return h.launch(jobs[0]); + var streams: [lanes.Engine.max_together]*lanes.Stream = undefined; + for (jobs[0..n], 0..) |job, i| streams[i] = &job.stream; + var errs: [lanes.Engine.max_together]?anyerror = undefined; + h.core.addStreams(streams[0..n], errs[0..n]); + for (jobs[0..n], errs[0..n]) |job, err| { + if (err) |e| { + _ = switch (e) { + error.Cancelled => h.cancel(job), + error.LaneMemory => h.requeue(job), + else => h.drop(job, h.words(e)), + }; + continue; + } + h.prefilled(job, job.began); + if (h.deliver(job)) h.remove(job); + } + return true; +} + +/// The next queued request admitted alone: its stream set up and launched. +fn one(h: *LaneHost) bool { + const job = h.takeJob() orelse return false; + if (!h.setup(job)) return true; + return h.launch(job); +} diff --git a/zig/src/core/lanes/backend.zig b/zig/src/core/lanes/backend.zig index fe26d3f003..23a4efd892 100644 --- a/zig/src/core/lanes/backend.zig +++ b/zig/src/core/lanes/backend.zig @@ -1,4 +1,4 @@ -//! What the round loop asks a backend (Metal, CUDA or the fake) to compute; drawn tokens stay on the GPU behind handles. +//! What the round loop asks a backend (Metal, CUDA or the fake); drawn tokens stay on the GPU behind handles. const std = @import("std"); const Stream = @import("stream.zig").Stream; const Shape = @import("shape.zig").Shape; @@ -68,6 +68,8 @@ pub const Backend = struct { prefill: *const fn (ptr: *anyopaque, s: *Stream) anyerror!void, /// One chunk of the prompt pass, false while chunks remain; null: the pass runs whole. prefill_step: ?*const fn (ptr: *anyopaque, s: *Stream) anyerror!bool = null, + /// Several streams' prompts in one pass, each as after `prefill`; false, doing nothing, when they don't fit. + prefill_many: ?*const fn (ptr: *anyopaque, ss: []const *Stream) anyerror!bool = null, /// Draw the first token from the prompt's last row at `position`; a handle to it. first: *const fn (ptr: *anyopaque, s: *Stream, position: u64) anyerror!u64, /// Feed one token and queue the draw of the next at `position` (a one-token round); a handle to it. @@ -137,7 +139,7 @@ pub const Backend = struct { b.vtable.release(b.ptr, s); } - /// A driver of its own starts a stream: the prompt pass and its first token, committed; a cancel or failure releases it. + /// A driver of its own starts a stream: prompt pass and first token, committed; cancel or failure releases it. pub fn opening(b: Backend, gpa: std.mem.Allocator, s: *Stream) !u32 { errdefer b.release(s); try b.prefill(s); diff --git a/zig/src/core/lanes/engine.zig b/zig/src/core/lanes/engine.zig index cf0a97ea4a..e4c7950781 100644 --- a/zig/src/core/lanes/engine.zig +++ b/zig/src/core/lanes/engine.zig @@ -83,7 +83,16 @@ pub const Engine = struct { pub const fills = open.fills; pub const fillStream = open.fillStream; - /// Take over a stream another driver decoded so far: its cache settled, `pending` and `cache_len` set, no drafts held. + pub const adoptPrefilled = open.adoptPrefilled; + pub const max_together = open.max_together; + pub const addStreams = open.addStreams; + + /// Whether the backend prefills several prompts in one pass (addStreams batches them). + pub fn batchesPrompts(e: *const Engine) bool { + return e.backend.vtable.prefill_many != null; + } + + /// Take over a stream another driver decoded: cache settled, `pending` and `cache_len` set, no drafts held. pub fn adopt(e: *Engine, s: *Stream) !void { try e.live.append(e.gpa, s); } diff --git a/zig/src/core/lanes/lanes.zig b/zig/src/core/lanes/lanes.zig index 5d50a16ef3..57d8ae9eb7 100644 --- a/zig/src/core/lanes/lanes.zig +++ b/zig/src/core/lanes/lanes.zig @@ -29,6 +29,7 @@ pub const grammar = @import("grammar.zig"); pub const Engine = engine.Engine; pub const Config = config.Config; +pub const Image = @import("stream.zig").Image; pub const Model = config.Model; pub const Stream = stream.Stream; pub const Sampling = sampling.Sampling; diff --git a/zig/src/core/lanes/open.zig b/zig/src/core/lanes/open.zig index 7cddf9a064..c75fa81b86 100644 --- a/zig/src/core/lanes/open.zig +++ b/zig/src/core/lanes/open.zig @@ -88,3 +88,58 @@ pub fn fillStream(e: *Engine, s: *Stream, first: bool) !bool { try opened(e, s); return true; } + +/// A stream another driver prefilled and drew `token` for: drafts asked, the token committed, rounds from here. +pub fn adoptPrefilled(e: *Engine, s: *Stream, token: u32) !void { + _ = e.arena.reset(.retain_capacity); + s.context.shrinkRetainingCapacity(s.prompt_len); + s.pending = null; + s.cache_len = s.prompt_len; + if (e.cfg.family_mtp and s.drafts) { + const d: u32 = @intCast(try e.rule.depth(win.who(s))); + try e.backend.draft(&.{.{ .stream = s, .follow = &.{}, .first = .{ .value = token }, .rows = null, .start = s.prompt_len, .position = s.prompt_len + 1, .depth = d }}); + s.dropHeld(e.gpa); + s.next = .{ .count = d }; + if (e.backend.vtable.tree) |tree| if (try tree(e.backend.ptr, s, e.gpa)) |held| { + s.next = held; + }; + } + if (s.logprobs != null) return error.LogprobsUnsupported; // another driver's prompt pass gives no first row + _ = try s.commit(e.gpa, &.{token}, &.{}); + s.pending = token; + if (s.finished) return e.release(s); + try e.live.append(e.gpa, s); +} + +/// Streams a backend prefills in one pass at most. +pub const max_together = 16; + +/// open.add for several streams, prompts in one pass when the backend allows; ``errs[i]``: stream i's error. +pub fn addStreams(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (errs) |*x| x.* = null; + const many = e.backend.vtable.prefill_many orelse return addEach(e, ss, errs); + if (ss.len < 2 or ss.len > max_together) return addEach(e, ss, errs); + _ = e.arena.reset(.retain_capacity); + for (ss) |s| trail.event(e, &.{ f("ev", str("add")), f("stream", str(s.id)) }) catch {}; + const together = many(e.backend.ptr, ss) catch |err| { + for (errs) |*x| x.* = err; + return; + }; + if (!together) return addEach(e, ss, errs); + for (ss, errs) |s, *x| { + if (s.isCancelled()) { // cancelled in its pass: no first token + e.backend.release(s); + x.* = error.Cancelled; + continue; + } + opened(e, s) catch |err| { + x.* = err; + }; + } +} + +fn addEach(e: *Engine, ss: []const *Stream, errs: []?anyerror) void { + for (ss, errs) |s, *x| add(e, s) catch |err| { + x.* = err; + }; +} diff --git a/zig/src/core/lanes/sampling.zig b/zig/src/core/lanes/sampling.zig index e9d469756c..43bddfa3a1 100644 --- a/zig/src/core/lanes/sampling.zig +++ b/zig/src/core/lanes/sampling.zig @@ -75,7 +75,7 @@ pub fn pairwiseSum(a: []const f64) f64 { return pairwiseSum(a[0..n2]) + pairwiseSum(a[n2..]); } -const Candidate = struct { value: f64, id: u64 }; +pub const Candidate = struct { value: f64, id: u64 }; fn byValueThenId(_: void, a: Candidate, b: Candidate) bool { return a.value > b.value or (a.value == b.value and a.id < b.id); @@ -88,6 +88,12 @@ pub fn choose(gpa: Allocator, values: []const f64, ids: []const u64, position: u defer gpa.free(order); for (order, values, ids) |*o, v, id| o.* = .{ .value = v, .id = id }; std.mem.sort(Candidate, order, {}, byValueThenId); + return chooseOrdered(gpa, order, position, s); +} + +/// ``choose`` over candidates already in its order (by value, then the lower id at a tie). +pub fn chooseOrdered(gpa: Allocator, order: []const Candidate, position: u64, s: Sampling) !u64 { + const width = order.len; if (s.temperature <= 0) return order[0].id; // temperature 0 is greedy: the first maximum, ties to the lower id const k = @max(1, @min(if (s.top_k != 0) @as(usize, s.top_k) else width, width)); const scaled = try gpa.alloc(f64, k); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index f028d10cc4..6d634c9600 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -76,7 +76,7 @@ pub const CancelCheck = struct { check: *const fn (ptr: *anyopaque) bool, }; -/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and where it keeps new ones. +/// A kept prompt state the backend restores before its prompt pass (core/prompt_cache.zig), and keeps new ones in. pub const Reuse = struct { saved: ?*anyopaque = null, // the backend's copy of the state after `at` prompt tokens (null: start at 0) at: u32 = 0, @@ -97,6 +97,9 @@ pub const Held = struct { parents: ?[]i32 = null, }; +/// An image for a vision tower: patches (fp32 [h * w, 1536], merge-block order), grid [1, h, w]; prompt placeholders. +pub const Image = struct { patches: []const f32, grid: [3]i64 }; + pub const Spec = struct { id: []const u8, prompt: []const u32, @@ -123,6 +126,7 @@ pub const Spec = struct { prompt_logprobs: ?u8 = null, // and for each prompt token after the first, with ranks (the prompt pass starts at 0) grammar: ?*grammar.Constraint = null, // structured output: the reply stays in it (the creator frees it) decode_spans: []const [2]u32 = &.{}, // prompt rows an earlier reply decoded (Request.decode_spans), for backends that keep their bits + images: []const Image = &.{}, // an image prompt's inputs, in prompt order (the backend encodes them at prefill) }; pub const Stream = struct { @@ -140,6 +144,7 @@ pub const Stream = struct { think_close: []const u32, think_end: i64, think_open: bool, + images: []const Image = &.{}, loop_guard: bool, loop_period: ?u32 = null, chunks: []const u32, @@ -202,6 +207,7 @@ pub const Stream = struct { .cancel_check = spec.cancel_check, .yield_check = spec.yield_check, .think_budget = spec.think_budget, + .images = spec.images, .think_close = spec.think_close, .think_end = spec.think_end, .think_open = spec.think_open orelse (spec.think_budget > 0 or (spec.loop_guard and spec.think_end >= 0)), diff --git a/zig/src/core/prompt_cache.zig b/zig/src/core/prompt_cache.zig index f1465325e1..81d5cf4ed4 100644 --- a/zig/src/core/prompt_cache.zig +++ b/zig/src/core/prompt_cache.zig @@ -108,6 +108,8 @@ pub const Entry = struct { last: []u32, // the last prompt that kept or resumed it shared: bool = false, // kept at a shared system block's cut: other conversations resume it, so its own never supersedes it branched: bool = false, // resumed again after a later prompt kept a longer state extending it: never superseded + pins: u32 = 0, // admitted jobs that restore it and have not yet: eviction passes it over + dead: bool = false, // removed while pinned: no lookup finds it, and its last unpin frees it }; pub const Counts = struct { hits: u64 = 0, misses: u64 = 0, kept: u64 = 0, evicted: u64 = 0, refused: u64 = 0, failed: u64 = 0 }; @@ -148,6 +150,7 @@ pub const Store = struct { } fn free(s: *Store, e: *Entry) void { + s.held -= e.bytes; s.family.vtable.drop(s.family.ptr, e.saved); s.gpa.free(e.tokens); s.gpa.free(e.decode_spans); @@ -157,8 +160,19 @@ pub const Store = struct { fn remove(s: *Store, i: usize) void { const e = s.entries.orderedRemove(i); - s.held -= e.bytes; - s.free(e); + e.dead = true; + if (e.pins == 0) s.free(e); + } + + /// `e` waits for a restore: nothing frees it until its last `unpin` (removed meanwhile, no lookup finds it). + pub fn pin(_: *Store, e: *Entry) *Entry { + e.pins += 1; + return e; + } + + pub fn unpin(s: *Store, e: *Entry) void { + e.pins -= 1; + if (e.pins == 0 and e.dead) s.free(e); } /// The family's free storage kept for later saves. @@ -450,7 +464,7 @@ pub const Store = struct { if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held -| bytes); // spare that cannot take this state goes before any state if (s.spare() < free_) continue; } - s.remove(s.victimBut(null, prompt, born).?); // its storage may come back spare and take this state + s.remove(s.victimBut(null, prompt, born) orelse return s.fail(at, error.RestoresPending)); // spare for this state s.counts.evicted += 1; } if (s.entries.items.len == 0) if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| bytes); @@ -473,11 +487,11 @@ pub const Store = struct { const charged = if (s.family.vtable.charged) |f| f(s.family.ptr, saved) else bytes; e.* = .{ .tokens = tokens, .decode_spans = decoded, .at = at, .saved = saved, .bytes = charged, .born = born, .used = s.clock, .last = last, .shared = shared }; owns_decoded = false; + s.held += charged; // free takes it back s.entries.append(s.gpa, e) catch { s.free(e); return s.fail(at, error.OutOfMemory); }; - s.held += charged; s.counts.kept += 1; if (shared) learned.learn(s, e, starts); if (s.family.vtable.trim) |f| f(s.family.ptr, s.budget -| s.held); // spare storage only inside what the budget leaves @@ -534,17 +548,17 @@ pub const Store = struct { return false; } - /// The entry to free first, never `skip`: a superseded one, no branch point, oldest; else the oldest, rewind last. + /// The entry to free first, never `skip` or pinned: superseded, no branch point, oldest; else oldest, rewind last. fn victimBut(s: *const Store, skip: ?*Entry, prompt: []const u32, born: u32) ?usize { var best: ?usize = null; for (s.entries.items, 0..) |e, i| { - if (e.shared or e.branched or e == skip or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used + if (e.shared or e.branched or e == skip or e.pins > 0 or s.isActiveRewind(e, prompt)) continue; // shared cuts and branch points go only as the least recently used const now = born > e.born and prompt.len > e.tokens.len and std.mem.eql(u32, prompt[0..e.tokens.len], e.tokens); // a longer request's tokens extend it if ((now or s.extended(e)) and (best == null or e.used < s.entries.items[best.?].used)) best = i; } if (best) |i| return i; var oldest: ?usize = null; - for (s.entries.items, 0..) |e, i| if (e != skip and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { + for (s.entries.items, 0..) |e, i| if (e != skip and e.pins == 0 and (oldest == null or s.before(e, s.entries.items[oldest.?], prompt))) { oldest = i; }; return oldest; diff --git a/zig/src/core/prompt_cache_test.zig b/zig/src/core/prompt_cache_test.zig index 8452036b75..cde079b745 100644 --- a/zig/src/core/prompt_cache_test.zig +++ b/zig/src/core/prompt_cache_test.zig @@ -391,3 +391,80 @@ test "a kept state grows only into what the budget leaves, and evicts nothing to var stranger: u8 = 0; try std.testing.expect(!s.grow(&stranger, 0)); // a state the store doesn't hold } + +test "an entry pinned for a restore still to come outlives the lookups set up after it; unpinned, it goes again" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 220); // fake bytes: 100 + at + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); // 103 bytes, the oldest + const y = [_]u32{ 6, 6, 6, 6 }; + _ = f.pass(&s, &y, try s.begin(a, &y, 3, &.{}, &.{}, null, &.{})); // 103: 206 held + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const l = try s.lookup(a, &xa, 3, &.{}, &.{}, &.{}); // a job set up to resume x + const e = l.entry.?; + _ = s.pin(e); + const z = [_]u32{ 7, 7, 7, 7, 7, 7, 7, 7 }; + _ = try s.lookup(a, &z, 7, &.{}, &.{}, &.{}); // the next job's room: y goes, not x + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + try std.testing.expect(s.find(&.{ 6, 6, 6, 6 }, &.{}, &.{}) == null); + const w: [130]u32 = @splat(8); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); // 220 bytes, and nothing but x to evict: it stays + try std.testing.expect(!s.keep(&w, 120, null, &.{}, &.{})); // and a keep with no room is refused + try std.testing.expect(s.find(&xa, &.{}, &.{}) == e); + s.unpin(e); + _ = try s.lookup(a, &w, 120, &.{}, &.{}, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(s.entries.items.len, f.live); +} + +test "two jobs pin one entry and the first restore fails: lookups miss it, and it lives until the second job unpins" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const xb = [_]u32{ 5, 5, 5, 5, 8 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + try std.testing.expect(s.pin((try s.lookup(a, &xb, 3, &.{}, &.{}, &.{})).entry.?) == e); + const held = s.held; + s.resumed(e, &xa, false); // the first job's restore failed: it prefilled from 0 + s.unpin(e); + try std.testing.expect(s.find(&xb, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 0), s.entries.items.len); + try std.testing.expectEqual(@as(usize, 1), f.live); // the second job's restore still reads it + try std.testing.expectEqual(held, s.held); + try std.testing.expectEqual(@as(u32, 1), e.pins); + s.resumed(e, &xb, true); // no longer kept: its report changes nothing + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); + try std.testing.expectEqual(@as(u64, 0), s.held); +} + +test "forget leaves a pinned entry to its restore" { + const gpa = std.testing.allocator; + var arena: std.heap.ArenaAllocator = .init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var f: Fake = .{ .gpa = gpa }; + var s = Store.init(gpa, f.snapshots(), .{ .min_prompt = 0 }, 1 << 20); + defer s.deinit(); + const x = [_]u32{ 5, 5, 5, 5 }; + _ = f.pass(&s, &x, try s.begin(a, &x, 3, &.{}, &.{}, null, &.{})); + const xa = [_]u32{ 5, 5, 5, 5, 9 }; + const e = s.pin((try s.lookup(a, &xa, 3, &.{}, &.{}, &.{})).entry.?); + s.forget(&xa, 3, &.{}); + try std.testing.expect(s.find(&xa, &.{}, &.{}) == null); + try std.testing.expectEqual(@as(usize, 1), f.live); + s.unpin(e); + try std.testing.expectEqual(@as(usize, 0), f.live); +} diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index c940202bea..e59b21daf2 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -6,6 +6,7 @@ const stream_ = @import("stream.zig"); const kernels = @import("kernels.zig"); const grouped = @import("grouped.zig"); const driver = @import("driver.zig"); +const ggml = @import("ggml.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -13,20 +14,20 @@ const Function = module.Function; const Stream = stream_.Stream; const Plan = grouped.Plan; -pub const Format = enum { affine4, fp8g, nvfp4 }; +pub const Format = enum { affine4, fp8g, nvfp4, iq2xxs }; /// One layer's experts on the device as each format's make() lays them out; a family stacks its shared ones last. pub const Layer = struct { format: Format, - up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144] - down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144] + up: u64, // affine4 [E][NI/32][D/64][288] int32; fp8g [E][NI/32][D/32][2][32][2][4]; nvfp4 [E][NI/32][D/32][2][144]; iq2xxs [E][NI/16][D/256][1056] bytes (ggml.zig) + down: u64, // affine4 [E][D/32][NI/64][288]; fp8g [E][D/32][NI/32][1][32][2][4]; nvfp4 [E][D/32][NI/32][1][144]; iq2xxs [E][D/16][NI/256][1056] up_scale: u64 = 0, // fp8g [E][2][NI/128][D/128] fp32; nvfp4 [E][2] global scales; affine4's are in the blocks down_scale: u64 = 0, // fp8g [E][1][D/128][NI/128]; nvfp4 [E][1] width: usize, // NI dims: usize, // D experts: usize, limit: f32 = 0.0, // fp8g, nvfp4: the SwiGLU clamp, 0 for none - relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up + relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up; iq2xxs: always }; /// A call's epilogue: up with its format's activation (affine4 relu^2, else SwiGLU), down to fp32 or to bf16 sums. @@ -40,6 +41,7 @@ pub const Grouped = struct { affine4: ?*const Affine4Experts = null, fp8g: ?*const Fp8Experts = null, nvfp4: ?*const Nvfp4Experts = null, + iq2xxs: ?*const Iq2xxsExperts = null, /// Decode rows: `items` bounds the plan's items (grouped.maxItems at tile 16); `skip` an expert left out, or -1. pub fn decode(g: Grouped, s: Stream, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { @@ -55,6 +57,7 @@ pub const Grouped = struct { }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, false, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).run(s, epi, in, l, p, out, items, skip), + .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, false, epi, in, l, p, out, items, skip), } } @@ -72,6 +75,7 @@ pub const Grouped = struct { }, .fp8g => return (g.fp8g orelse return error.FormatNotLoaded).run(s, epi, in, l, p, true, out, items, skip), .nvfp4 => return (g.nvfp4 orelse return error.FormatNotLoaded).runPrompt(s, epi, in, l, p, out, items, skip), + .iq2xxs => return iq2Run(g.iq2xxs orelse return error.FormatNotLoaded, s, true, epi, in, l, p, out, items, skip), } } }; @@ -79,10 +83,26 @@ pub const Grouped = struct { /// Pairs a prompt item holds, per format (the Python engines' prompt tiles). pub fn promptTile(f: Format) usize { return switch (f) { - .affine4, .fp8g, .nvfp4 => 64, + .affine4, .fp8g, .nvfp4, .iq2xxs => 64, }; } +/// IQ2_XXS experts (ggml.cu): relu^2 up, down to fp32 or bf16; packed by Iq2xxsExperts.pack from the raw blocks. +pub const Iq2xxsExperts = ggml.Kernels; + +fn iq2Run(k: *const Iq2xxsExperts, s: Stream, prompt: bool, epi: Epi, in: Rows, l: Layer, p: Plan, out: u64, items: usize, skip: i32) !void { + if (l.limit != 0) return error.OptionNotBuilt; + const up = epi == .up; + const o: ggml.Out = switch (epi) { + .up => .relu2, + .down_f32 => .f32, + .down_bf16 => .bf16, + }; + const n = if (up) l.width else l.dims; + const kk = if (up) l.dims else l.width; + return k.experts(s, prompt, .iq2xxs, o, in.x, in.stride, if (up) in.slots else 0, if (up) l.up else l.down, n, kk, p, out, items, skip); +} + fn int(x: usize) c_int { return @intCast(x); } diff --git a/zig/src/cuda/ggml.zig b/zig/src/cuda/ggml.zig new file mode 100644 index 0000000000..e7bb967ac3 --- /dev/null +++ b/zig/src/cuda/ggml.zig @@ -0,0 +1,217 @@ +//! Bolt's ggml blocks (ggml.cu): IQ2_XXS and Q8_0 repacked into 16-column tiles, for grouped experts and dense rows. +const std = @import("std"); +const driver = @import("driver.zig"); +const module = @import("module.zig"); +const launch_ = @import("launch.zig"); +const stream_ = @import("stream.zig"); +const kernels = @import("kernels.zig"); +const grouped = @import("grouped.zig"); + +const Driver = driver.Driver; +const Module = module.Module; +const Function = module.Function; +const Stream = stream_.Stream; + +/// The checkpoint's block formats: IQ2_XXS (66 bytes a 256 inputs), Q8_0 (34 bytes a 32). +pub const Kind = enum { iq2xxs, q8_0 }; + +/// Inputs and bytes of one checkpoint block. +pub fn block(kind: Kind) struct { k: usize, bytes: usize } { + return switch (kind) { + .iq2xxs => .{ .k = 256, .bytes = 66 }, + .q8_0 => .{ .k = 32, .bytes = 34 }, + }; +} + +/// Bytes of 16 columns' blocks side by side (a packed tile's step along K). +pub fn tileBytes(kind: Kind) usize { + return 16 * block(kind).bytes; +} + +/// Bytes of an (n, k) matrix, the checkpoint's and the packed layout's alike. +pub fn bytes(kind: Kind, n: usize, k: usize) usize { + const b = block(kind); + return n * (k / b.k) * b.bytes; +} + +/// One packed projection on the device: n outputs from k inputs. +pub const Weight = struct { kind: Kind, w: u64, n: usize, k: usize }; + +pub fn q8Weight(w: u64, n: usize, k: usize) Weight { + return .{ .kind = .q8_0, .w = w, .n = n, .k = k }; +} + +pub fn iq2Weight(w: u64, n: usize, k: usize) Weight { + return .{ .kind = .iq2xxs, .w = w, .n = n, .k = k }; +} + +/// What a call writes: fp32, bf16 of relu(x)^2 (Nemotron's expert activation), or bf16. +pub const Out = enum { f32, relu2, bf16 }; + +pub const prompt_rows = 64; // rows a prompt item holds (the plan's prompt tile) +pub const decode_rows = 16; // rows a decode item holds at most (the plan's decode tile) +const prompt_warps = 8; // 16-column tiles a prompt CTA +/// Decode CTAs: 4 warps split a tile's K, or (short K: fewer than `split_loads` loads a tile) a warp a tile, 4 tiles. +const split_loads = 16; +const prompt_smem: u32 = 3 * 64 * 128 * 2; // three stages of 64 rows' 128 inputs + +/// Where byte `o` of a packed matrix comes from in the checkpoint's [n][k / KB][B] (tf_ggml_pack's map). +pub fn source(kind: Kind, n: usize, k: usize, o: usize) usize { + _ = n; + const b = block(kind); + const kb = k / b.k; + const tb_bytes = tileBytes(kind); + const tb = o / tb_bytes; + const r = o % tb_bytes; + const tile = tb / kb; + const blk = tb % kb; + var col: usize = undefined; + var at: usize = undefined; + if (r >= tb_bytes - 32) { + const q = r - (tb_bytes - 32); + col = tile * 16 + (q >> 2) + 8 * ((q >> 1) & 1); + at = q & 1; + } else if (kind == .iq2xxs) { + col = tile * 16 + ((r >> 4) & 7) + 8 * ((r >> 3) & 1); + at = 2 + (r >> 7) * 8 + (r & 7); + } else { + const ln = r >> 4; + const q = (r >> 2) & 3; + col = tile * 16 + (ln >> 2) + 8 * (q & 1); + at = 2 + 8 * (ln & 3) + 4 * (q >> 1) + (r & 3); + } + return (col * kb + blk) * b.bytes + at; +} + +/// The host repack of one (n, k) matrix, as tf_ggml_pack does it on the device. +pub fn packHost(kind: Kind, out: []u8, raw: []const u8, n: usize, k: usize) void { + std.debug.assert(out.len == bytes(kind, n, k) and raw.len == out.len and n % 16 == 0); + for (out, 0..) |*o, i| o.* = raw[source(kind, n, k, i)]; +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +/// The ggml module's kernels: decode (GEMV) and prompt entries per format and output, and the repack. +pub const Kernels = struct { + mod: Module, + gemv: [2][2][3]Function, // by Kind, K split or 4 tiles a CTA, then Out + prompt: [2][3]Function, + pack_fn: Function, + + /// Entry names: `form` is "gemv" (K split), "gemv4" (4 tiles a CTA) or "prompt". + pub fn symbol(comptime kind: Kind, comptime form: []const u8, comptime out: Out) [:0]const u8 { + return "tf_ggml_" ++ (if (kind == .iq2xxs) "iq2" else "q8") ++ "_" ++ form ++ "_" ++ @tagName(out); + } + + pub fn load(d: *const Driver) !Kernels { + if (!kernels.available) return error.BuiltWithoutKernels; + var mod = try Module.load(d, kernels.ggml); + errdefer mod.unload(); + var g: Kernels = .{ .mod = mod, .gemv = undefined, .prompt = undefined, .pack_fn = try mod.function("tf_ggml_pack") }; + inline for (.{ Kind.iq2xxs, Kind.q8_0 }, 0..) |kind, ki| inline for (.{ Out.f32, Out.relu2, Out.bf16 }, 0..) |out, oi| { + g.gemv[ki][0][oi] = try mod.function(symbol(kind, "gemv", out)); + g.gemv[ki][1][oi] = try mod.function(symbol(kind, "gemv4", out)); + g.prompt[ki][oi] = try mod.function(symbol(kind, "prompt", out)); + try g.prompt[ki][oi].allowDynamicShared(prompt_smem); + }; + return g; + } + + pub fn unload(g: *Kernels) void { + g.mod.unload(); + } + + /// `count` checkpoint matrices [n][k / KB][B] back to back at `raw` into packed tiles at `out` (as many bytes). + pub fn pack(g: *const Kernels, s: Stream, kind: Kind, raw: u64, out: u64, n: usize, k: usize, count: usize) !void { + if (n % 16 != 0 or k % block(kind).k != 0) return error.ShapeNotTiled; + const total = count * bytes(kind, n, k); + if (total == 0) return; + var a: launch_.Args = .{}; + a.add(raw); + a.add(out); + a.add(@as(c_int, @intFromBool(kind == .iq2xxs))); + a.add(int(n)); + a.add(int(k / block(kind).k)); + a.add(@as(i64, @intCast(total))); + try launch_.launch(g.pack_fn, .{ .grid = .{ .x = @intCast(@min((total + 255) / 256, 1 << 20)), .y = 1, .z = 1 }, .block = .{ .x = 256 } }, s, &a); + } + + fn check(kind: Kind, n: usize, k: usize, prompt: bool) !void { + if (n % 16 != 0 or k % block(kind).k != 0 or (prompt and k % 128 != 0)) return error.ShapeNotTiled; + } + + /// The common launch: `items`, `counts` and `members` 0 for dense rows (`rows` of them), else a plan's. + fn go(g: *const Kernels, s: Stream, prompt: bool, out_kind: Out, kind: Kind, x: u64, stride: usize, slots: usize, w: u64, n: usize, k: usize, p: ?grouped.Plan, out: u64, units: usize, rows: usize, skip: i32) !void { + try check(kind, n, k, prompt); + const ki = @backingInt(kind); + const oi = @backingInt(out_kind); + const nt16 = n / 16; + if (units == 0) return; + var a: launch_.Args = .{}; + a.add(x); + a.add(int(stride)); + a.add(int(slots)); + a.add(w); + a.add(int(k / 32)); + a.add(int(nt16)); + for ([_]u64{ if (p) |q| q.items else 0, if (p) |q| q.counts else 0, if (p) |q| q.members else 0, out }) |v| a.add(v); + a.add(int(n)); + a.add(int(rows)); + a.add(skip); + if (prompt) { + const grid = units * ((nt16 + prompt_warps - 1) / prompt_warps); + const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 32 * prompt_warps }, .shared = prompt_smem }; + return launch_.launch(g.prompt[ki][oi], cfg, s, &a); + } + const loads = k / 32 / @as(usize, if (kind == .iq2xxs) 4 else 1); + const tiles: usize = if (loads < split_loads) 4 else 1; + const grid = (units * nt16 + tiles - 1) / tiles; // items past the plan's count return at once + try launch_.launch(g.gemv[ki][@intFromBool(tiles > 1)][oi], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128 } }, s, &a); + } + + /// Dense decode rows (at most 16): out (rows, n) = x (rows, k; rows `ldx` apart, 0 packed) times w. + pub fn decode(g: *const Kernels, s: Stream, w: Weight, x: u64, ldx: usize, rows: usize, out: u64, out_kind: Out) !void { + if (rows > decode_rows) return error.TooManyRows; + return g.go(s, false, out_kind, w.kind, x, if (ldx == 0) w.k else ldx, 0, w.w, w.n, w.k, null, out, @intFromBool(rows > 0), rows, -1); + } + + /// Dense prompt rows (any count): 64 rows by 128 columns a CTA. + pub fn promptRows(g: *const Kernels, s: Stream, w: Weight, x: u64, ldx: usize, rows: usize, out: u64, out_kind: Out) !void { + return g.go(s, true, out_kind, w.kind, x, if (ldx == 0) w.k else ldx, 0, w.w, w.n, w.k, null, out, (rows + prompt_rows - 1) / prompt_rows, rows, -1); + } + + /// Either form by row count: decode at 16 rows or fewer, else prompt. + pub fn matmul(g: *const Kernels, s: Stream, w: Weight, x: u64, ldx: usize, rows: usize, out: u64, out_kind: Out) !void { + return if (rows <= decode_rows) g.decode(s, w, x, ldx, rows, out, out_kind) else g.promptRows(s, w, x, ldx, rows, out, out_kind); + } + + /// Grouped experts, at most `items` items: [E] packed (n, k) at `w`; x token rows if `slots` > 0, else pair rows. + pub fn experts(g: *const Kernels, s: Stream, prompt: bool, kind: Kind, out_kind: Out, x: u64, stride: usize, slots: usize, w: u64, n: usize, k: usize, p: grouped.Plan, out: u64, items: usize, skip: i32) !void { + return g.go(s, prompt, out_kind, kind, x, stride, slots, w, n, k, p, out, items, 0, skip); + } +}; + +test "the repack is a permutation of each matrix's bytes" { + const gpa = std.testing.allocator; + inline for (.{ Kind.iq2xxs, Kind.q8_0 }) |kind| { + const n = 32; + const k = 512; + const len = bytes(kind, n, k); + const seen = try gpa.alloc(bool, len); + defer gpa.free(seen); + @memset(seen, false); + for (0..len) |o| { + const src = source(kind, n, k, o); + try std.testing.expect(!seen[src]); + seen[src] = true; + } + } + // q8_0: tile 1, block 3, lane 13 (gq 3, t 1), step 1, column gq + 8, byte 2: column 27, input 8 + 4 + 2 + const o = (1 * 16 + 3) * 544 + 13 * 16 + 3 * 4 + 2; + try std.testing.expectEqual((27 * 16 + 3) * 34 + 2 + 8 + 4 + 2, source(.q8_0, 32, 512, o)); + // iq2xxs: tile 0, block 1, group 5, gq 6, column gq + 8, byte 7; then column 6's d byte 1 + try std.testing.expectEqual((14 * 2 + 1) * 66 + 2 + 40 + 7, source(.iq2xxs, 32, 512, 1056 + 5 * 128 + 6 * 16 + 8 + 7)); + try std.testing.expectEqual((6 * 2 + 1) * 66 + 1, source(.iq2xxs, 32, 512, 1056 + 1024 + 6 * 4 + 1)); +} diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 0fa8bf6e28..d8f77b0dfd 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -33,12 +33,14 @@ pub const train: []const u8 = if (available) &Blob("fatbin_train").bytes else &. pub const qmmf: []const u8 = if (available) &Blob("fatbin_qmmf").bytes else &.{}; pub const prompt16: []const u8 = if (available) &Blob("fatbin_prompt16").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; +pub const ggml: []const u8 = if (available) &Blob("fatbin_ggml").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; pub const torch_indexing: []const u8 = if (available) &Blob("fatbin_torch_indexing").bytes else &.{}; pub const torch_movement: []const u8 = if (available) &Blob("fatbin_torch_movement").bytes else &.{}; pub const torch_nemotron_constants: []const u8 = if (available) &Blob("fatbin_torch_nemotron_constants").bytes else &.{}; +pub const radio_vision: []const u8 = if (available) &Blob("fatbin_radio_vision").bytes else &.{}; pub const nemotron_norms: []const u8 = if (available) &Blob("fatbin_nemotron_norms").bytes else &.{}; pub const nemotron_route: []const u8 = if (available) &Blob("fatbin_nemotron_route").bytes else &.{}; pub const nemotron_mamba: []const u8 = if (available) &Blob("fatbin_nemotron_mamba").bytes else &.{}; diff --git a/zig/src/cuda/qlinear.zig b/zig/src/cuda/qlinear.zig index 6d6a56f5a3..1f7d42510e 100644 --- a/zig/src/cuda/qlinear.zig +++ b/zig/src/cuda/qlinear.zig @@ -7,6 +7,7 @@ const memory = @import("memory.zig"); const stream_ = @import("stream.zig"); const qmmf = @import("qmmf.zig"); const nvfp4 = @import("nvfp4.zig"); +const ggml = @import("ggml.zig"); const Driver = driver.Driver; const Module = module.Module; @@ -36,11 +37,14 @@ pub const Weight = union(enum) { nvfp4: qmmf.Weight, fp8: qmmf.Weight, // e4m3 with one fp32 scale (ModelOpt FP8) bf16: qmmf.Weight, // unquantized: the checkpoint's bf16 values + q8_0: ggml.Weight, // ggml Q8_0 packed by ggml.Kernels.pack + iq2xxs: ggml.Weight, // ggml IQ2_XXS packed by ggml.Kernels.pack pub fn outputs(w: Weight) usize { return switch (w) { .affine4 => |q| q.n, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.n, + .q8_0, .iq2xxs => |q| q.n, }; } @@ -48,6 +52,7 @@ pub const Weight = union(enum) { return switch (w) { .affine4 => |q| q.k, .fp8g, .nvfp4, .fp8, .bf16 => |q| q.k, + .q8_0, .iq2xxs => |q| q.k, }; } }; @@ -60,6 +65,7 @@ pub const Linear = struct { affine4: ?*const Affine4Kernels = null, lane: ?*const qmmf.Lane = null, // decode rows of every lane format; fp8g and bf16 prompt rows too nvfp4_prompt: ?*const nvfp4.Prompt = null, // the prompt GEMM: nvfp4 and fp8 prompt rows (linear.py's prefill) + ggml: ?*const ggml.Kernels = null, // q8_0 and iq2xxs rows of both kinds /// Decode rows: out (rows, n) bf16, each row's bits the same at every row count the format takes. pub fn decode(l: Linear, s: Stream, w: Weight, in: Rows, out: u64, rows: usize) !void { @@ -73,6 +79,7 @@ pub const Linear = struct { .nvfp4 => |q| return l.onLane(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onLane(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + .q8_0, .iq2xxs => |q| return (l.ggml orelse return error.FormatNotLoaded).decode(s, q, in.x, in.ldx, rows, out, .bf16), } } @@ -88,6 +95,8 @@ pub const Linear = struct { .nvfp4 => |q| return l.onPrompt(s, q, .fp4, in, out, rows), .fp8 => |q| return l.onPrompt(s, q, .fp8, in, out, rows), .bf16 => |q| return l.onLane(s, q, .bf16, in, out, rows), + // prompt rows always take the prompt GEMM: a row's bits must not depend on where a prompt's chunks are cut + .q8_0, .iq2xxs => |q| return (l.ggml orelse return error.FormatNotLoaded).promptRows(s, q, in.x, in.ldx, rows, out, .bf16), } } diff --git a/zig/src/cuda/root.zig b/zig/src/cuda/root.zig index 3053c09a6b..5c1a4932e6 100644 --- a/zig/src/cuda/root.zig +++ b/zig/src/cuda/root.zig @@ -36,6 +36,7 @@ pub const embed_rows = @import("embed_rows.zig"); pub const experts = @import("experts.zig"); pub const sample = @import("sample.zig"); pub const modelopt = @import("modelopt.zig"); +pub const ggml = @import("ggml.zig"); test { _ = @import("memory.zig"); @@ -54,4 +55,5 @@ test { _ = carveout; _ = @import("nvfp4.zig"); _ = @import("modelopt.zig"); + _ = @import("ggml.zig"); } diff --git a/zig/src/families/nemotron/config.zig b/zig/src/families/nemotron/config.zig index aec6e207f0..5eb9192282 100644 --- a/zig/src/families/nemotron/config.zig +++ b/zig/src/families/nemotron/config.zig @@ -3,7 +3,7 @@ const std = @import("std"); pub const Kind = enum { mamba, moe, attention }; -pub const max_layers = 64; +pub const max_layers = 96; // Nemotron 3.5 Bolt has 88 /// The MTP head's acceptance at depth 1, 2, ... given the ones before it, until a stream has its own (model.py draft_prior). pub const draft_prior = [_]f64{ 0.8, 0.72, 0.68, 0.62, 0.58, 0.55, 0.5, 0.5 }; @@ -31,6 +31,9 @@ pub const pocket_key = "tensorfold_learned_neurons"; /// How the checkpoint stores its matrices: MLX affine 4-bit, or NVIDIA ModelOpt (FP8, NVFP4 or bf16, read per tensor). pub const Format = enum { mlx, modelopt }; +/// One MoE layer's routed experts and their width. +pub const Moe = struct { experts: usize = 0, width: usize = 0 }; + pub const Config = struct { hidden: usize, vocab: usize, @@ -44,9 +47,12 @@ pub const Config = struct { heads: usize, kv_heads: usize, head_dim: usize, - experts: usize, + experts: usize, // the widest MoE layer's routed experts (every layer's, unless `moe` says otherwise) top_k: usize, - expert_width: usize, + expert_width: usize, // the widest MoE layer's routed expert width + moe: [max_layers]Moe = @splat(.{}), // each MoE layer's own experts and width (block_configs); zero: the uniform ones + omni: bool = false, // nemotron_h_omni: the language model's tensors under language_model., a vision tower beside it + latent: usize = 0, // moe_latent_size: routed experts read and write this width through fc1/fc2_latent_proj (LatentMoE) shared_width: usize, // the shared expert's own neurons, learned ones (pocket) left out pocket: usize = 0, // tensorfold_learned_neurons: learned neurons after the shared expert's own, in the same tensors routed_scaling: f32, @@ -60,6 +66,17 @@ pub const Config = struct { eos: [4]u32 = .{ 0, 0, 0, 0 }, eos_count: usize = 0, + /// MoE layer `i`'s routed experts and width. + pub fn layerMoe(self: Config, i: usize) Moe { + const m = self.moe[i]; + return if (m.experts == 0) .{ .experts = self.experts, .width = self.expert_width } else m; + } + + /// The width routed experts read and write: the latent one, else hidden. + pub fn expertDims(self: Config) usize { + return if (self.latent > 0) self.latent else self.hidden; + } + pub fn inner(self: Config) usize { return self.mamba_heads * self.mamba_head_dim; } @@ -76,9 +93,9 @@ pub const Config = struct { return (self.heads + 2 * self.kv_heads) * self.head_dim; } - /// A token's expert pairs: its routed experts, then the shared expert's two halves. + /// A token's expert pairs: its routed experts, then the shared expert's two halves (LatentMoE runs it on its own). pub fn slots(self: Config) usize { - return self.top_k + 2; + return if (self.latent > 0) self.top_k else self.top_k + 2; } pub fn count(self: Config, kind: Kind) usize { @@ -96,7 +113,7 @@ pub const Config = struct { pub fn read(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Config { const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); defer gpa.free(path); - const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 22)); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 25)); defer gpa.free(text); var why: Why = .{}; return parse(gpa, text, &why) catch |e| { @@ -118,6 +135,8 @@ fn number(v: std.json.Value) !f64 { return switch (v) { .float => |f| f, .integer => |i| @floatFromInt(i), + // transformers 5 writes non-finite floats as {"__float__": "Infinity"} + .object => |o| if (o.get("__float__")) |f| (if (f == .string) std.fmt.parseFloat(f64, f.string) catch error.BadConfig else error.BadConfig) else error.BadConfig, else => error.BadConfig, }; } @@ -225,8 +244,11 @@ fn learned(path: []const u8) bool { pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config { var parsed = try std.json.parseFromSlice(std.json.Value, allocator, text, .{}); defer parsed.deinit(); - const o = parsed.value.object; - if (!std.mem.eql(u8, (o.get("model_type") orelse return error.BadConfig).string, "nemotron_h")) return error.NotNemotronH; + const root = parsed.value.object; + // nemotron_h_omni (Bolt) wraps the language model's config in llm_config, its quantization at the top + const wrapped = isType(root, "nemotron_h_omni"); + const o = if (wrapped) (llmConfig(root) orelse return error.BadConfig) else root; + if (!isType(o, "nemotron_h") and !isType(o, "nemotron_h_puzzle")) return error.NotNemotronH; var c = Config{ .hidden = try int(o, "hidden_size"), .vocab = try int(o, "vocab_size"), @@ -239,9 +261,10 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config .heads = try int(o, "num_attention_heads"), .kv_heads = try int(o, "num_key_value_heads"), .head_dim = 0, - .experts = try int(o, "n_routed_experts"), - .top_k = try int(o, "num_experts_per_tok"), - .expert_width = try int(o, "moe_intermediate_size"), + // a puzzle config (Bolt) names these per block + .experts = if (o.get("block_configs") != null) 0 else try int(o, "n_routed_experts"), + .top_k = if (o.get("block_configs") != null) 0 else try int(o, "num_experts_per_tok"), + .expert_width = if (o.get("block_configs") != null) 0 else try int(o, "moe_intermediate_size"), .shared_width = try int(o, "moe_shared_expert_intermediate_size"), .routed_scaling = if (o.get("routed_scaling_factor")) |v| @floatCast(try number(v)) else 1.0, .eps = if (o.get("layer_norm_epsilon")) |v| @floatCast(try number(v)) else 1e-5, @@ -253,8 +276,13 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config c.shared_width -= c.pocket; } if (o.get("norm_topk_prob")) |v| c.norm_topk = v.bool; + c.omni = wrapped; + if (o.get("moe_latent_size")) |v| if (v == .integer) { + c.latent = @intCast(v.integer); + }; + if (o.get("block_configs")) |v| if (v == .array and v.array.items.len > 0) try puzzle(&c, v.array.items, why); // hybrid_override_pattern (one character a block) wins over layers_block_type, as the Python loader reads them - if (o.get("hybrid_override_pattern")) |v| if (v == .string and v.string.len > 0) { + if (c.layers == 0) if (o.get("hybrid_override_pattern")) |v| if (v == .string and v.string.len > 0) { if (v.string.len > max_layers) return error.BadConfig; for (v.string, 0..) |ch, i| c.kinds[i] = try kindOf(&.{ch}); c.layers = v.string.len; @@ -265,25 +293,66 @@ pub fn parse(allocator: std.mem.Allocator, text: []const u8, why: *Why) !Config for (blocks, 0..) |b, i| c.kinds[i] = try kindOf(b.string); c.layers = blocks.len; } - if (o.get("eos_token_id")) |e| switch (e) { - .integer => |i| { - c.eos[0] = @intCast(i); - c.eos_count = 1; - }, - .array => |a| for (a.items[0..@min(a.items.len, 4)]) |x| { - c.eos[c.eos_count] = @intCast(x.integer); - c.eos_count += 1; - }, - else => {}, - }; + // generation_config's [2, 11] is the omni wrapper's eos plus the language model's + if (wrapped) if (root.get("eos_token_id")) |e| eos(&c, e); + if (o.get("eos_token_id")) |e| eos(&c, e); if (o.get("time_step_limit")) |v| if (v == .array and v.array.items.len == 2) { c.dt_min = @floatCast(try number(v.array.items[0])); c.dt_max = @floatCast(try number(v.array.items[1])); }; - try quantization(&c, o, why); + try quantization(&c, root, why); return c; } +fn isType(o: std.json.ObjectMap, t: []const u8) bool { + const v = o.get("model_type") orelse return false; + return v == .string and std.mem.eql(u8, v.string, t); +} + +fn llmConfig(o: std.json.ObjectMap) ?std.json.ObjectMap { + const v = o.get("llm_config") orelse return null; + return if (v == .object) v.object else null; +} + +/// An eos id, or a list of them, added once each. +fn eos(c: *Config, e: std.json.Value) void { + const items: []const std.json.Value = switch (e) { + .integer => &.{e}, + .array => |a| a.items, + else => &.{}, + }; + for (items) |x| if (x == .integer and !c.isEos(@intCast(x.integer)) and c.eos_count < c.eos.len) { + c.eos[c.eos_count] = @intCast(x.integer); + c.eos_count += 1; + }; +} + +/// nemotron_h_puzzle's block_configs: a block type per layer, each MoE layer its own experts and width (one top-k). +fn puzzle(c: *Config, blocks: []const std.json.Value, why: *Why) !void { + if (blocks.len > max_layers) { + why.set("config.json has {d} layers; the native Nemotron engine reads up to {d}", .{ blocks.len, max_layers }); + return error.BadConfig; + } + c.experts = 0; + c.expert_width = 0; + for (blocks, 0..) |b, i| { + if (b != .object) return error.BadConfig; + c.kinds[i] = try kindOf((b.object.get("block_type") orelse return error.BadConfig).string); + if (c.kinds[i] != .moe) continue; + const m: Moe = .{ .experts = try int(b.object, "n_routed_experts"), .width = try int(b.object, "moe_intermediate_size") }; + const k = try int(b.object, "num_experts_per_tok"); + if (c.top_k == 0) c.top_k = k; + if (k != c.top_k) { + why.set("layer {d} routes to a different number of experts than the others; the engine reads one top-k", .{i}); + return error.BadConfig; + } + c.moe[i] = m; + c.experts = @max(c.experts, m.experts); + c.expert_width = @max(c.expert_width, m.width); + } + c.layers = blocks.len; +} + /// The shapes our kernels serve (Nemotron 3.5 Lightning 30B-A3B, 4-bit groups of 64): Metal's generated sources, CUDA's captured cubins. pub fn checkShapes(c: Config) !void { const ok = c.hidden == 2688 and c.vocab == 131072 and c.mamba_heads == 64 and c.mamba_head_dim == 64 and @@ -372,3 +441,46 @@ test "parse the Lightning config" { try std.testing.expectEqual(std.math.inf(f32), c.dt_max); try checkShapes(c); } + +test "a Bolt config: omni wrapper, puzzle blocks with their own experts and widths, LatentMoE" { + const bolt = + \\{"model_type": "nemotron_h_omni", "eos_token_id": 11, "llm_config": {"model_type": "nemotron_h_puzzle", + \\ "hidden_size": 4096, "vocab_size": 131072, "mamba_num_heads": 128, "mamba_head_dim": 64, "n_groups": 8, + \\ "ssm_state_size": 128, "conv_kernel": 4, "num_attention_heads": 32, "num_key_value_heads": 2, "head_dim": 128, + \\ "num_experts": 512, "moe_latent_size": 1024, "moe_shared_expert_intermediate_size": 5376, + \\ "routed_scaling_factor": 5.0, "eos_token_id": 2, "time_step_limit": [0.0, {"__float__": "Infinity"}], + \\ "block_configs": [{"block_type": "mamba"}, + \\ {"block_type": "moe", "moe_intermediate_size": 2048, "n_routed_experts": 240, "num_experts_per_tok": 9}, + \\ {"block_type": "attention"}, + \\ {"block_type": "moe", "moe_intermediate_size": 2560, "n_routed_experts": 496, "num_experts_per_tok": 9}]}, + \\ "quantization_config": {"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION"}} + ; + var why: Why = .{}; + const c = try parse(std.testing.allocator, bolt, &why); + try std.testing.expectEqual(@as(usize, 4), c.layers); + try std.testing.expectEqual(Kind.attention, c.kinds[2]); + try std.testing.expectEqual(@as(usize, 9), c.top_k); + try std.testing.expectEqual(@as(usize, 9), c.slots()); + try std.testing.expectEqual(@as(usize, 1024), c.expertDims()); + try std.testing.expectEqual(Moe{ .experts = 240, .width = 2048 }, c.layerMoe(1)); + try std.testing.expectEqual(Moe{ .experts = 496, .width = 2560 }, c.layerMoe(3)); + try std.testing.expectEqual(@as(usize, 496), c.experts); + try std.testing.expectEqual(@as(usize, 2560), c.expert_width); + try std.testing.expectEqual(@as(usize, 18560), c.projDim()); + try std.testing.expect(c.isEos(2) and c.isEos(11)); + try std.testing.expectEqual(std.math.inf(f32), c.dt_max); + try std.testing.expectEqual(Format.modelopt, c.format); + try std.testing.expect(c.omni); +} + +test "the Bolt checkpoint's own config.json, when it is on this machine" { + const io = std.testing.io; + const dir = "/storage/models/nemotron-bolt-74b-24gb"; + std.Io.Dir.cwd().access(io, dir ++ "/config.json", .{}) catch return error.SkipZigTest; + const c = try Config.read(std.testing.allocator, io, dir); + try std.testing.expectEqual(@as(usize, 88), c.layers); + try std.testing.expectEqual(@as(usize, 40), c.count(.moe)); + try std.testing.expectEqual(@as(usize, 8), c.count(.attention)); + try std.testing.expectEqual(@as(usize, 496), c.experts); + try std.testing.expectEqual(Moe{ .experts = 240, .width = 2048 }, c.layerMoe(1)); +} diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index cb2fa934a1..a91bf0b8ca 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -18,6 +18,8 @@ pub const Head = @import("cuda_mtp.zig").Head; pub const Lanes = @import("cuda_lanes.zig").Cuda; pub const native = @import("cuda_native.zig"); pub const train = @import("cuda_train.zig"); +/// Bolt's image tower (C-RADIO) and projector: one image's patches to the rows its `` placeholders take. +pub const vision = @import("cuda_vision.zig"); pub const sites = @import("cuda_sites.zig"); pub const train_ops = @import("cuda_train_ops.zig"); pub const glue = @import("cuda_glue.zig"); @@ -27,6 +29,7 @@ pub const reuse = @import("cuda_reuse.zig"); pub const slide = @import("cuda_slide.zig"); test { + _ = vision; _ = @import("config.zig"); _ = @import("draft_ids.zig"); _ = kernels; diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..c93f0b21ce 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -123,6 +123,9 @@ pub const Engine = struct { armed: marks.Armed = .{}, // the marks inside the prompt chunk about to run (cuda_marks.zig) spares: marks.Spares = .{}, // free Mamba buffers the prompt cache set aside for the next marks last_chunk: usize = 0, // where the last prompt pass's last chunk started (the prompt's last row: len - 1 - this) + snap_host: bool = false, // kept prompt states in pinned host memory (cuda_reuse.zig), set when the host opens the engine + features: u64 = 0, // during an image prompt's prefill: the tower's rows, spliced over its placeholders (Bolt) + feature_src: u64 = 0, // and each prompt row's feature row, -1 for text /// Loads the checkpoint into the Python engine's layouts and sizes the caches; `triton_dir` null: our own glue. pub fn init(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, model_dir: []const u8, triton_dir: ?[]const u8, opts: Options) !*Engine { @@ -151,12 +154,18 @@ pub const Engine = struct { if (opts.slide) return error.SlideNeedsMlx; // the learner reads MLX 4-bit projections try e.k.loadModelopt(ctx); } + // Bolt (LatentMoE, ggml blocks): no MTP head in the checkpoint; our own glue (captured sets are Lightning's) + const mtp = opts.mtp and e.c.latent == 0; + if (e.c.latent > 0) { + if (e.k.triton != null) return error.CapturedGlueIsLightnings; + try e.k.loadGgml(ctx); + } const chunks = e.max_len / state.chunk_keys; e.nch = if (e.k.triton) |*t| @intCast(t.smallestConst("_chunk", "NCH", @intCast(chunks)) orelse { std.log.err("no captured attention kernel covers {d} chunks of 512 keys; capture one for this context", .{chunks}); return error.ContextUnsupported; }) else chunks; - e.w = try weights.load(gpa, io, e.ops(), model_dir, e.c, opts.mtp); + e.w = try weights.load(gpa, io, e.ops(), model_dir, e.c, mtp); errdefer e.w.deinit(); e.b = try state.Buffers.init(ctx.d, e.c, e.max_len, e.nch, e.carve); errdefer e.b.deinit(); @@ -175,7 +184,7 @@ pub const Engine = struct { e.sampled_ready = try cuda.Event.init(ctx.d, false); try e.copied.record(e.stream); try e.setSampling(opts.sampling); - if (opts.graphs) try e.capture(opts.mtp); + if (opts.graphs) try e.capture(mtp); e.heat_gate = heat.Gate.fromEnv() catch |err| { std.log.err("TF_HEAT_HIGH and TF_HEAT_LOW: both or neither, in degrees, with low at or under high", .{}); return err; @@ -342,6 +351,8 @@ pub const Engine = struct { f.dump = dump; f.slot = e.slot; f.untaught = e.bound.untaught; + f.features = e.features; + f.feature_src = e.feature_src; if (dump == null and e.c.experts <= 128 and e.c.slots() <= 8) f.side = e.side; // tf_plan_routed's bounds return f; } diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index a0e0452e86..65467cb1e1 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -60,6 +60,8 @@ pub const Forward = struct { untaught: bool = false, // the pocket leaves rows (the model as shipped): `segs`' untaught sequences', else all segs: []const Seg = &.{}, // a window's sequences, as windowSegs runs them captures: []const kept.Mark = &.{}, // marks inside the prompt chunk: their Mamba states are copied out as it runs + features: u64 = 0, // an image prompt's tower rows (bf16 [image rows, hidden]), 0 for a text prompt + feature_src: u64 = 0, // and each prompt row's feature row (i32 [prompt], -1 for text) pub fn init(c: Config, w: *const weights.Weights, b: *const state.Buffers, ops: kern.Ops, max_len: usize, nch: usize, sampled: bool) Forward { return .{ @@ -212,7 +214,7 @@ pub const Forward = struct { }, .moe => { try f.experts(blk.moe, rows, false); - delta = .moe; + delta = if (blk.moe.latent != null) .dense else .moe; }, } } @@ -245,8 +247,9 @@ pub const Forward = struct { const o = f.ops; const ex = m.experts; const pairs = rows * c.slots(); + if (m.latent) |lt| return f.latentExperts(m, lt, rows, prompt); if (!prompt) if (f.side) |sd| return f.forked(m, rows, sd); - try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.routed_scaling, c.norm_topk); + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); try f.mark(.route); const tile: usize = if (prompt) cuda.experts.promptTile(ex.format) else 16; try o.plan(b.pick, pairs, ex.experts, tile, b.plan); @@ -267,6 +270,52 @@ pub const Forward = struct { try f.sharedRest(o, m, rows, !prompt); } + /// LatentMoE (Bolt), into b.delta: fc2_latent(sum of the routed experts on fc1_latent(y)) + the shared expert on y. + fn latentExperts(f: *const Forward, m: weights.MoE, lt: weights.Latent, rows: usize, prompt: bool) !void { + const c = f.c; + const b = f.b; + const o = f.ops; + const ex = m.experts; + const K = c.top_k; + const pairs = rows * K; + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, ex.experts, K, K, c.routed_scaling, c.norm_topk); + try f.mark(.route); + if (prompt) try o.prefillDense(b.y, lt.fc1, b.lat, rows) else try o.dense(b.y, b.xs, lt.fc1, b.lat, rows); + const tile: usize = if (prompt) cuda.experts.promptTile(ex.format) else 16; + const g = o.grouped(); + try o.plan(b.pick, pairs, ex.experts, tile, b.plan); + try f.mark(.plan); + var items = kern.maxItems(pairs, ex.experts, tile); + const up_in: cuda.experts.Rows = .{ .x = b.lat, .stride = c.latent, .slots = K }; + const down_in: cuda.experts.Rows = .{ .x = b.act, .stride = ex.width }; + if (prompt) { + try g.prompt(o.s, .up, up_in, ex, b.plan, b.act, items, -1); + try g.prompt(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1); + } else { + try g.decode(o.s, .up, up_in, ex, b.plan, b.act, items, -1); + try f.mark(.experts_up); + try g.decode(o.s, .down_f32, down_in, ex, b.plan, b.ymoe, items, -1); + try f.mark(.experts_down); + } + try o.latentSum(b.ymoe, b.wts, b.z, rows, K, c.latent); + if (prompt) try o.prefillDense(b.z, lt.fc2, b.delta, rows) else try o.dense(b.z, 0, lt.fc2, b.delta, rows); + // the shared expert: every row's pair with expert 0 of a one-expert layer, relu^2 between as the routed ones + const sh = lt.shared; + try o.plan(b.zeros, rows, 1, tile, b.plan); + items = kern.maxItems(rows, 1, tile); + const sh_up: cuda.experts.Rows = .{ .x = b.y, .stride = c.hidden, .slots = 1 }; + const sh_down: cuda.experts.Rows = .{ .x = b.act, .stride = sh.width }; + if (prompt) { + try g.prompt(o.s, .up, sh_up, sh, b.plan, b.act, items, -1); + try g.prompt(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1); + } else { + try g.decode(o.s, .up, sh_up, sh, b.plan, b.act, items, -1); + try g.decode(o.s, .down_f32, sh_down, sh, b.plan, b.ymoe, items, -1); + } + try o.addBf16(b.delta, b.delta, b.ymoe, true, rows * c.hidden); + try f.mark(.shared_wait); + } + /// The pocket over the rows of sequences that read what was learned: every row, or the runs between untaught ones'. fn pocket(f: *const Forward, o: kern.Ops, p: weights.Pocket, d: usize, y_f32: bool, rows: usize) !void { if (!f.untaught) return f.pocketRun(o, p, d, y_f32, .{ 0, rows }); @@ -312,7 +361,7 @@ pub const Forward = struct { try g.decode(sd.s, .down_f32, in[1], ex, sp, b.ymoe, 2, -1); try f.sharedRest(so, m, rows, true); try sd.join.record(sd.s); - try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.routed_scaling, c.norm_topk); + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); try f.mark(.route); try o.planRouted(b.pick, rows, c.slots(), c.top_k, c.experts, 16, b.plan); try f.mark(.plan); @@ -341,6 +390,7 @@ pub const Forward = struct { pub fn chunkBegin(f: *const Forward, w: *Walk) !void { const b = f.b; try f.embed(b.p_ids, b.emb, w.rows); + if (f.features != 0) try f.ops.spliceRows(b.emb, f.features, f.feature_src + @as(u64, w.pos) * 4, w.rows, f.c.hidden); w.x = b.emb; } @@ -427,7 +477,7 @@ pub const Forward = struct { }, .moe => { try f.experts(blk.moe, w.rows, true); - w.delta = .moe; + w.delta = if (blk.moe.latent != null) .dense else .moe; }, } } diff --git a/zig/src/families/nemotron/cuda_glue.zig b/zig/src/families/nemotron/cuda_glue.zig index 7d3f0c5ef5..4a3d63d347 100644 --- a/zig/src/families/nemotron/cuda_glue.zig +++ b/zig/src/families/nemotron/cuda_glue.zig @@ -88,10 +88,18 @@ const AttnSegs = extern struct { kc: [max_segs]u64, vc: [max_segs]u64, meta: [ma /// The router's K slices and the shared bytes its block stages: 16 rows and 16 experts of one slice, padded a word. pub fn routerShape(d: usize) struct { sk: usize, shared: u32 } { - const sk = triton.divisor(d / 64, 6); - return .{ .sk = sk, .shared = @intCast(2 * 16 * (d / sk / 2 + 1) * 4) }; + var sk = triton.divisor(d / 64, 6); + while (staged(d, sk) > 48 << 10 and (d / 64) % (sk * 2) == 0) sk *= 2; // wide rows (Bolt's): within 48 KiB, no opt-in + return .{ .sk = sk, .shared = @intCast(staged(d, sk)) }; } +fn staged(d: usize, sk: usize) usize { + return 2 * 16 * (d / sk / 2 + 1) * 4; +} + +/// The most K slices routerShape takes for a row of `d`: the router's scratch is that many [rows, E] fp32 sums. +pub const max_router_slices = 8; + /// Every launch of a window or prompt chunk's glue, on one stream. pub const Glue = struct { set: ?*const cuda.aot.Set, // a captured Triton set (GB10's qualified one), else our own kernels @@ -151,10 +159,10 @@ pub const Glue = struct { try g.go(g.f.add_moe_norm, .{ rows, 1, 1 }, 256, 0, &a); } - /// fp32 router logits over K slices, then each row's top-k and the shared halves' slots. - pub fn route(g: Glue, x: u64, w: u64, bias: u64, part: u64, ids: u64, wts: u64, rows: usize, d: usize, e: usize, top_k: usize, scaling: f32, norm: bool) !void { + /// fp32 router logits over K slices, then each row's top-k and the shared halves' slots (`slots` past top-k). + pub fn route(g: Glue, x: u64, w: u64, bias: u64, part: u64, ids: u64, wts: u64, rows: usize, d: usize, e: usize, top_k: usize, slots: usize, scaling: f32, norm: bool) !void { if (g.tri()) |t| return t.route(x, w, bias, part, ids, wts, rows, d, e, top_k, scaling, norm); - if (e > 256 or top_k + 2 > 8) return error.RouterTooWide; + if (e > 512 or top_k > 16 or slots < top_k or slots > top_k + 2) return error.RouterTooWide; // the top-k kernel's MAXC const shape = routerShape(d); var a: cuda.Args = .{}; ptrs(&a, &.{ x, w, part }); @@ -164,7 +172,7 @@ pub const Glue = struct { ptrs(&b, &.{ part, bias, ids, wts }); b.add(int(rows)); b.add(scaling); - ints(&b, &.{ e, shape.sk, top_k, top_k + 2 }); + ints(&b, &.{ e, shape.sk, top_k, slots }); b.add(@as(c_int, @intFromBool(norm))); try g.go(g.f.topk, .{ rows, 1, 1 }, 32, 0, &b); } @@ -353,4 +361,8 @@ test "the router stages one slice of 16 rows and 16 experts" { const s = routerShape(2688); try std.testing.expectEqual(@as(usize, 6), s.sk); try std.testing.expectEqual(@as(u32, 2 * 16 * 225 * 4), s.shared); + // Bolt's 4096-wide rows: four slices would stage 66 KiB, eight stay under 48 + const b = routerShape(4096); + try std.testing.expectEqual(@as(usize, 8), b.sk); + try std.testing.expect(b.shared <= 48 << 10 and b.sk <= max_router_slices); } diff --git a/zig/src/families/nemotron/cuda_kernels.zig b/zig/src/families/nemotron/cuda_kernels.zig index 7932f1c437..77e041bcee 100644 --- a/zig/src/families/nemotron/cuda_kernels.zig +++ b/zig/src/families/nemotron/cuda_kernels.zig @@ -61,11 +61,15 @@ pub const Kernels = struct { pocket_in: [2]cuda.Function, // the pocket's two halves a block row a row, then tiled for a prompt chunk's rows pocket_out: [2]cuda.Function, add_row: cuda.Function, + latent_sum: cuda.Function, // LatentMoE: a row's routed latent outputs weighted and summed (Bolt) + add_bf16: cuda.Function, + splice_rows: cuda.Function, // an image prompt's rows take the vision tower's features (Bolt) draws: cuda.sample.Draws, // sample.cu: the keyed draws and logprob rows grammar_mask: cuda.Function, // structured output: disallowed tokens -inf before the draw torch: torch_ops.Functions, train: train_ops.Kernels, // Sliding Weights' learning step, and the learned change a forward adds while lessons run modelopt: ?Modelopt = null, // loaded by loadModelopt for a ModelOpt checkpoint, else never + ggml: ?cuda.ggml.Kernels = null, // loaded by loadGgml for a checkpoint with ggml blocks (Bolt's IQ2_XXS and Q8_0) gb10: bool, discrete: bool, // the card has its own memory: checkpoint bytes reach it through page-locked slots @@ -93,6 +97,9 @@ pub const Kernels = struct { k.pocket_in = .{ try k.mods[7].function("tf_pocket_in"), try k.mods[7].function("tf_pocket_in_rows") }; k.pocket_out = .{ try k.mods[7].function("tf_pocket_out"), try k.mods[7].function("tf_pocket_out_rows") }; k.add_row = try k.mods[7].function("tf_add_row"); + k.latent_sum = try k.mods[7].function("tf_latent_sum"); + k.add_bf16 = try k.mods[7].function("tf_add_bf16"); + k.splice_rows = try k.mods[7].function("tf_splice_rows"); k.torch = try torch_ops.Functions.resolve(k.mods[8..14]); k.draws = try cuda.sample.Draws.resolve(k.mods[14]); k.glue = try glue.Fns.resolve(k.mods[16..21]); @@ -102,6 +109,7 @@ pub const Kernels = struct { errdefer k.norms.unload(); k.grammar_mask = try k.mods[14].function("tf_grammar_mask"); k.modelopt = null; + k.ggml = null; k.triton = if (triton_dir) |dir| try cuda.aot.Set.load(gpa, io, d, ctx.device, dir) else null; errdefer if (k.triton) |*t| t.deinit(); k.train = try train_ops.Kernels.load(d); @@ -126,6 +134,7 @@ pub const Kernels = struct { m.experts.unload(); m.pack.unload(); } + if (k.ggml) |*g| g.unload(); k.affine.deinit(); k.train.deinit(); if (k.triton) |*t| t.deinit(); @@ -146,6 +155,13 @@ pub const Kernels = struct { errdefer experts.unload(); k.modelopt = .{ .lane = lane, .gemm = gemm, .experts = experts, .pack = try cuda.modelopt.Packer.load(k.d) }; } + + /// The ggml block kernels (IQ2_XXS experts, IQ2_XXS and Q8_0 projections, their repack) Bolt's export needs. + pub fn loadGgml(k: *Kernels, ctx: *const cuda.Context) !void { + if (k.ggml != null) return; + _ = ctx; + k.ggml = try cuda.ggml.Kernels.load(k.d); + } }; /// qmm.split_k (cuda/qlinear.zig). @@ -204,12 +220,44 @@ pub const Ops = struct { /// Every projection format the kernels loaded (cuda/qlinear.zig). pub fn linear(o: Ops) cuda.qlinear.Linear { const m = if (o.k.modelopt) |*x| x else null; - return .{ .affine4 = &o.k.affine, .lane = if (m) |x| &x.lane else null, .nvfp4_prompt = if (m) |x| &x.gemm else null }; + return .{ .affine4 = &o.k.affine, .lane = if (m) |x| &x.lane else null, .nvfp4_prompt = if (m) |x| &x.gemm else null, .ggml = if (o.k.ggml) |*g| g else null }; } /// Every grouped-expert format the kernels loaded (cuda/experts.zig). pub fn grouped(o: Ops) cuda.experts.Grouped { - return .{ .affine4 = &o.k.experts, .nvfp4 = if (o.k.modelopt) |*x| &x.experts else null }; + return .{ .affine4 = &o.k.experts, .nvfp4 = if (o.k.modelopt) |*x| &x.experts else null, .iq2xxs = if (o.k.ggml) |*g| g else null }; + } + + /// LatentMoE: z[r] = sum_k wt[r, k] y[r * k_slots + k], rows of `l` bf16 from fp32 pair rows (tf_latent_sum). + pub fn latentSum(o: Ops, y: u64, wt: u64, z: u64, rows: usize, k_slots: usize, l: usize) !void { + var a: cuda.Args = .{}; + a.add(y); + a.add(wt); + a.add(z); + a.add(int(k_slots)); + a.add(int(l)); + try o.go(o.k.latent_sum, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// Rows of h whose src (i32, one a row) names a feature row take that row of `features` ([*, d] bf16). + pub fn spliceRows(o: Ops, h: u64, features: u64, src: u64, rows: usize, d: usize) !void { + var a: cuda.Args = .{}; + a.add(h); + a.add(features); + a.add(src); + a.add(int(d)); + try o.go(o.k.splice_rows, .{ rows, 1, 1 }, 256, 0, &a); + } + + /// out = a + b, n bf16 values (b bf16, or fp32 when `b_f32`). + pub fn addBf16(o: Ops, out: u64, a_: u64, b: u64, b_f32: bool, n: usize) !void { + var a: cuda.Args = .{}; + a.add(out); + a.add(a_); + a.add(b); + a.add(@as(c_int, @intFromBool(b_f32))); + a.add(int(n)); + try o.go(o.k.add_bf16, .{ (n + 255) / 256, 1, 1 }, 256, 0, &a); } /// Decode rows of a projection on this stream (the affine-4 paths' bits unchanged); xs: x's 64-group sums. diff --git a/zig/src/families/nemotron/cuda_lanes.zig b/zig/src/families/nemotron/cuda_lanes.zig index 203ccd74cb..9146fca387 100644 --- a/zig/src/families/nemotron/cuda_lanes.zig +++ b/zig/src/families/nemotron/cuda_lanes.zig @@ -1,6 +1,7 @@ //! Nemotron CUDA streams own their sequences and share window rounds. const std = @import("std"); +const vision = @import("cuda_vision.zig"); const cuda = @import("cuda"); const lanes = @import("lanes"); const core = @import("core"); @@ -48,6 +49,7 @@ pub const Cuda = struct { measured: ?core.draft_depth.Costs = null, // a lone stream's depth rule prices its rounds by these bits: ?cuda.HostBuffer = null, // structured output's row bits, device-mapped (row_slots rows), made at first use bits_dev: u64 = 0, + tower: ?*vision.Tower = null, // --vision (Bolt): an image prompt's images are encoded at its prefill pub fn init(gpa: std.mem.Allocator, e: *Engine, head: ?*Head) !Cuda { var pinned = try cuda.HostBuffer.alloc(e.ctx.d, engine.max_streams * state.max_rows * 4); @@ -63,6 +65,7 @@ pub const Cuda = struct { var it = self.lanes.valueIterator(); while (it.next()) |l| if (!l.own) self.e.freeSeq(l.seq); self.lanes.deinit(self.gpa); + if (self.tower) |t| t.deinit(); self.pinned.free(); self.rows_out.free(); if (self.bits) |*b| b.free(); @@ -254,8 +257,21 @@ pub const Cuda = struct { e.decide = if (s.decide != null and e.lessons() != null) s.decide.? else lesson_route.none; try e.setSampling(s.sampling); s.cached = 0; - const from = if (s.prompt_logprobs != null) 0 else self.restoreKept(s, ids.len); // prompt rows need every row - const keep: ?engine.Keep = if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; + // an image prompt: tower rows spliced over the placeholders per chunk; never resumed or kept (same ids) + var media: ?vision.Media = null; + if (s.images.len > 0) { + const tower = self.tower orelse return error.NemotronServedWithoutVision; + media = try vision.Media.init(self.gpa, tower, e.stream, ids, s.images); + e.features = media.?.features.ptr; + e.feature_src = media.?.src.ptr; + } + defer if (media) |*m| { + e.features = 0; + e.feature_src = 0; + m.deinit(); + }; + const from = if (s.prompt_logprobs != null or media != null) 0 else self.restoreKept(s, ids.len); // prompt rows need every row + const keep: ?engine.Keep = if (s.reuse.marks.len == 0 or media != null) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; var sink: PromptSink = .{ .gpa = self.gpa, .s = s }; s.prompt_rows.clearRetainingCapacity(); if (s.prompt_logprobs) |k| { diff --git a/zig/src/families/nemotron/cuda_modelopt.zig b/zig/src/families/nemotron/cuda_modelopt.zig index 546b861b39..5055bbe165 100644 --- a/zig/src/families/nemotron/cuda_modelopt.zig +++ b/zig/src/families/nemotron/cuda_modelopt.zig @@ -13,9 +13,24 @@ const Loader = W.Loader; const Tensor = core.checkpoint.Tensor; const Nvfp4Experts = cuda.experts.Nvfp4Experts; -const Kind = enum { nvfp4, fp8, bf16 }; +const Kind = enum { nvfp4, fp8, bf16, q8_0, iq2xxs }; + +/// Bytes and weights of one ggml block: Q8_0 34 per 32, IQ2_XXS 66 per 256. +fn ggmlBlock(k: Kind) [2]usize { + return switch (k) { + .q8_0 => .{ 34, 32 }, + .iq2xxs => .{ 66, 256 }, + else => unreachable, + }; +} fn kindOf(t: Tensor) !Kind { + // ggml blocks (Bolt): [rows, blocks, block bytes]; NVFP4 packs pairs as [rows, k / 2] + if (t.dtype == .u8 and t.rank == 3) return switch (t.dim(2)) { + 34 => .q8_0, + 66 => .iq2xxs, + else => error.UnsupportedQuantization, + }; return switch (t.dtype) { .u8 => .nvfp4, .f8_e4m3 => .fp8, @@ -42,16 +57,22 @@ fn scalar(ck: *core.Checkpoint, n: []const u8) !f32 { fn part(ck: *core.Checkpoint, pre: []const u8) !Part { var buf: [192]u8 = undefined; const w = try ck.get(try name(&buf, "{s}.weight", .{pre})); - if (w.rank != 2) return error.UnexpectedTensor; + if (w.rank != 2 and w.rank != 3) return error.UnexpectedTensor; // we keep bf16 activations and a bf16 KV cache, more precise than the W8A8 and FP8-cache math these scales are for for ([_][]const u8{ "input_scale", "k_scale", "v_scale" }) |extra| { const n = try name(&buf, "{s}.{s}", .{ pre, extra }); if (ck.has(n)) _ = try ck.get(n); } const kind = try kindOf(w); - var p: Part = .{ .w = w, .kind = kind, .n = w.dim(0), .k = if (kind == .nvfp4) w.dim(1) * 2 else w.dim(1) }; + if ((w.rank == 3) != (kind == .q8_0 or kind == .iq2xxs)) return error.UnexpectedTensor; + const k = switch (kind) { + .nvfp4 => w.dim(1) * 2, + .q8_0, .iq2xxs => w.dim(1) * ggmlBlock(kind)[1], + else => w.dim(1), + }; + var p: Part = .{ .w = w, .kind = kind, .n = w.dim(0), .k = k }; switch (kind) { - .bf16 => {}, + .bf16, .q8_0, .iq2xxs => {}, .fp8 => p.scale = try scalar(ck, try name(&buf, "{s}.weight_scale", .{pre})), .nvfp4 => { p.blocks = try ck.expect(try name(&buf, "{s}.weight_scale", .{pre}), .f8_e4m3, &.{ p.n, p.k / 16 }); @@ -81,6 +102,7 @@ pub fn dense(L: *Loader, ck: *core.Checkpoint, label: []const u8, prefixes: []co var buf: [128]u8 = undefined; const s = L.ops.s; switch (first.kind) { + .q8_0, .iq2xxs => return ggmlDense(L, label, parts[0..prefixes.len], n, k), .bf16, .fp8 => { const size: usize = if (first.kind == .bf16) 2 else 1; const base = try L.tmp(n * k * size); @@ -115,9 +137,96 @@ pub fn dense(L: *Loader, ck: *core.Checkpoint, label: []const u8, prefixes: []co } } +/// Row-stacked ggml parts (q, k, v: rows after rows) uploaded as stored, then tiled on the GPU (a byte permutation). +fn ggmlDense(L: *Loader, label: []const u8, parts: []const Part, n: usize, k: usize) !kern.QLinear { + const g = try ggmlKernels(L); + const kind = parts[0].kind; + const base = try L.tmp(n * k / ggmlBlock(kind)[1] * ggmlBlock(kind)[0]); + var at: usize = 0; + for (parts) |p| { + try L.src.upload(base + at, p.w.bytes); + at += p.w.bytes.len; + } + var buf: [128]u8 = undefined; + const gk: cuda.ggml.Kind = if (kind == .q8_0) .q8_0 else .iq2xxs; + const out = try L.alloc(try name(&buf, "{s}.weight", .{label}), cuda.ggml.bytes(gk, n, k)); + try L.src.flush(); + try g.pack(L.ops.s, gk, base, out, n, k, 1); + return if (kind == .q8_0) .{ .q8_0 = cuda.ggml.q8Weight(out, n, k) } else .{ .iq2xxs = cuda.ggml.iq2Weight(out, n, k) }; +} + +fn ggmlKernels(L: *Loader) !*const cuda.ggml.Kernels { + return if (L.ops.k.ggml) |*g| g else error.GgmlNotLoaded; +} + +/// IQ2_XXS grouped experts from `count` tensors named by `names`, [n, k] each, uploaded back to back and tiled. +fn iq2Experts(L: *Loader, ck: *core.Checkpoint, label: []const u8, count: usize, n: usize, k: usize, names: anytype) !u64 { + const g = try ggmlKernels(L); + const one = n * k / 256 * 66; + const base = try L.tmp(count * one); + var b: [192]u8 = undefined; + for (0..count) |i| { + const p = try part(ck, try names.at(&b, i)); + if (p.kind != .iq2xxs or p.n != n or p.k != k) { + std.log.err("{s}: {t} [{d}, {d}], expected IQ2_XXS [{d}, {d}]", .{ try names.at(&b, i), p.kind, p.n, p.k, n, k }); + return error.UnexpectedTensor; + } + try L.src.upload(base + i * one, p.w.bytes); + } + const out = try L.alloc(label, cuda.ggml.bytes(.iq2xxs, n, k) * count); + try L.src.flush(); + try g.pack(L.ops.s, .iq2xxs, base, out, n, k, count); + return out; +} + +/// A LatentMoE layer (Bolt): router and bias, fc1/fc2_latent_proj, latent IQ2_XXS experts, full-width shared expert. +pub fn latent(L: *Loader, ck: *core.Checkpoint, label: []const u8, pre: []const u8, c: Config, layer: usize) !W.MoE { + const lm = c.layerMoe(layer); + const E = lm.experts; + const Lw = c.latent; + var a: [128]u8 = undefined; + var b: [192]u8 = undefined; + const r = try routerOf(L, try name(&a, "{s}.router", .{label}), try ck.get(try name(&b, "{s}gate.weight", .{pre})), E, c.hidden); + const bias = try ck.expect(try name(&b, "{s}gate.e_score_correction_bias", .{pre}), .f32, &.{E}); + var m: W.MoE = .{ .router = r, .bias = try L.raw(try name(&a, "{s}.bias", .{label}), bias), .experts = undefined }; + var ex: kern.Experts = .{ .format = .iq2xxs, .up = 0, .down = 0, .width = lm.width, .dims = Lw, .experts = E, .relu2 = true }; + const Names = struct { + pre: []const u8, + leaf: []const u8, + fn at(x: @This(), buf: []u8, i: usize) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}experts.{d}.{s}", .{ x.pre, i, x.leaf }); + } + }; + ex.up = try iq2Experts(L, ck, try name(&a, "{s}.experts.up", .{label}), E, lm.width, Lw, Names{ .pre = pre, .leaf = "up_proj" }); + ex.down = try iq2Experts(L, ck, try name(&a, "{s}.experts.down", .{label}), E, Lw, lm.width, Names{ .pre = pre, .leaf = "down_proj" }); + m.experts = ex; + const Shared = struct { + pre: []const u8, + leaf: []const u8, + fn at(x: @This(), buf: []u8, _: usize) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}shared_experts.{s}", .{ x.pre, x.leaf }); + } + }; + var sh: kern.Experts = .{ .format = .iq2xxs, .up = 0, .down = 0, .width = c.shared_width, .dims = c.hidden, .experts = 1, .relu2 = true }; + sh.up = try iq2Experts(L, ck, try name(&a, "{s}.shared.up", .{label}), 1, c.shared_width, c.hidden, Shared{ .pre = pre, .leaf = "up_proj" }); + sh.down = try iq2Experts(L, ck, try name(&a, "{s}.shared.down", .{label}), 1, c.hidden, c.shared_width, Shared{ .pre = pre, .leaf = "down_proj" }); + var p1: [192]u8 = undefined; + var p2: [192]u8 = undefined; + m.latent = .{ + .fc1 = try dense(L, ck, try name(&a, "{s}.fc1_latent", .{label}), &.{try name(&p1, "{s}fc1_latent_proj", .{pre})}), + .fc2 = try dense(L, ck, try name(&b, "{s}.fc2_latent", .{label}), &.{try name(&p2, "{s}fc2_latent_proj", .{pre})}), + .shared = sh, + }; + return m; +} + /// The router as the kernels read it: bf16 [E, D]; an fp32 router is taken only when every value is a bf16. fn router(L: *Loader, label: []const u8, t: Tensor, c: Config) !u64 { - if (t.rank != 2 or t.dim(0) != c.experts or t.dim(1) != c.hidden) return error.UnexpectedTensor; + return routerOf(L, label, t, c.experts, c.hidden); +} + +fn routerOf(L: *Loader, label: []const u8, t: Tensor, experts_: usize, hidden: usize) !u64 { + if (t.rank != 2 or t.dim(0) != experts_ or t.dim(1) != hidden) return error.UnexpectedTensor; if (t.dtype == .bf16) return L.raw(label, t); if (t.dtype != .f32) return error.UnexpectedTensor; const raw = try L.gpa.alloc(u8, t.bytes.len); diff --git a/zig/src/families/nemotron/cuda_native.zig b/zig/src/families/nemotron/cuda_native.zig index 64594342c4..4469f9151e 100644 --- a/zig/src/families/nemotron/cuda_native.zig +++ b/zig/src/families/nemotron/cuda_native.zig @@ -20,12 +20,18 @@ const core = @import("core"); pub const slide = @import("cuda_slide.zig"); pub const model_type = "nemotron_h"; -pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-nvfp4" }; +/// Nemotron 3.5 Bolt's config: the language model under llm_config, beside its vision tower's. +pub const also_model_types: []const []const u8 = &.{"nemotron_h_omni"}; +pub const formats: []const []const u8 = &.{ "mlx-q4g64", "modelopt-nvfp4", "modelopt-ggml" }; pub const default_context: i64 = engine.default_context; pub const max_segments: u32 = @import("cuda_segments.zig").MAX; pub const prompt_rows: u32 = state.pass_rows; -pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, vision: bool = false }; + +/// --vision on a Bolt checkpoint: its C-RADIO tower, images prepared and placed as qwen_image.nemotron lays them. +pub const vision_nemotron = true; +const vision = @import("cuda_vision.zig"); /// Thinking budgets the lone driver takes (cuda_lone.zig): above it, the stream goes to the lane core before its cut. pub const lone_min_budget = @import("cuda_lone.zig").min_budget; @@ -51,6 +57,7 @@ pub const Loaded = struct { target: *reuse.Target, // the prompt cache copies this engine (cuda_reuse.zig) learns: ?*engine.Engine = null, // --slide: the engine its learner trains, whose sites it attached at load learn_bytes: usize = 0, // --slide: what its trainer allocates at the first lesson, kept out of the streams' room + host_snaps: bool = false, // kept prompt states in pinned host memory (a discrete card): the cache's budget is the host's }; /// The first shard, index or config.json in `dir` that a write in place would reach through a link, or null. @@ -71,19 +78,27 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ errdefer if (carve) |c| closeCarveout(gpa, c); const e = try engine.Engine.init(gpa, io, ctx, dir, kernels, .{ .context = o.context, .mtp = o.drafts, .graphs = true, .sampling = null, .segments = o.segments, .carveout = carve, .slide = o.slide }); errdefer e.deinit(); - const head: ?*Head = if (o.drafts) try Head.init(e) else null; + // Bolt's checkpoint has no MTP head: it decodes serially whatever --drafts says + const drafts = o.drafts and e.c.latent == 0; + const head: ?*Head = if (drafts) try Head.init(e) else null; errdefer if (head) |h| h.deinit(); if (head) |h| try h.capture(); - if (!o.drafts) try e.captureWindows(); // the one-row window: Engine.init captures windows only when drafting + if (!drafts) try e.captureWindows(); // the one-row window: Engine.init captures windows only when drafting // the sampled mode's sets: a placeholder rule now, each stream's own uploaded at its prefill try e.setSampling(.{ .seed = 0, .temperature = 1.0 }); try e.captureWindows(); if (head) |h| try h.capture(); try e.setSampling(null); + e.snap_host = e.k.discrete; // a card with its own memory keeps prompt states in the host's (GB10 shares one) const own = try gpa.create(Owned); errdefer gpa.destroy(own); own.* = .{ .gpa = gpa, .e = e, .head = head, .lanes = try Lanes.init(gpa, e, head), .target = .{ .e = e, .head = head }, .carve = carve }; errdefer own.lanes.deinit(); + // --vision: the tower and its scratch for the largest image, loaded beside the model (Bolt alone has one) + if (o.vision) { + if (!e.c.omni) return error.NemotronCheckpointHasNoVision; + own.lanes.tower = try vision.Tower.load(gpa, io, ctx.d, dir); + } own.living = try slide.loadLiving(gpa, io, e, dir); // before the lanes' costs are measured: they decode through it errdefer if (own.living) |*s| s.deinit(); try own.lanes.measure(io, dir); @@ -98,6 +113,7 @@ pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [ .lone = if (head != null) loneRun else null, .learns = if (o.slide) e else null, .learn_bytes = if (o.slide) @import("cuda_train.zig").Trainer.memory(e.c) else 0, + .host_snaps = e.snap_host, }; } @@ -117,6 +133,9 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.PromptTooLong => "the prompt and its reply exceed this server's context window: shorten it or lower max_tokens", error.OutOfDeviceMemory => "the GPU had no memory left for this request's caches: retry once another request ends", + error.NemotronServedWithoutVision => "image inputs need the server started with --vision", + error.BadImagePlaceholders => "the image placeholders do not match the prepared images", + error.BadImageGrid, error.BadImagePatches => "an image's patches do not fit the vision tower", else => null, }; } diff --git a/zig/src/families/nemotron/cuda_prompt.zig b/zig/src/families/nemotron/cuda_prompt.zig index ba5408db8d..7765bab806 100644 --- a/zig/src/families/nemotron/cuda_prompt.zig +++ b/zig/src/families/nemotron/cuda_prompt.zig @@ -39,7 +39,7 @@ pub fn prefillFrom(e: *Engine, prompt: []const u32, from: usize, dump: ?*Dump, h if (head) |h| try h.reset(); } else if (e.pos != from) return error.BadReuse; const waited = e.heat_gate.waited_s; - if (from == 0 and keep == null and e.segments > 1 and dump == null and e.prompt_rows == null and prompt.len > state.prefill_rows and e.decide == lesson_route.none and !Cancel.background(cancel)) { + if (from == 0 and keep == null and e.features == 0 and e.segments > 1 and dump == null and e.prompt_rows == null and prompt.len > state.prefill_rows and e.decide == lesson_route.none and !Cancel.background(cancel)) { try segs.prefill(e, try e.segmentSet(), prompt, head, e.segments, cancel); } else try e.serialChunks(prompt, from, dump, head, cancel, keep); if (e.heat_gate.waited_s > waited) heat.note(e.heat_gate.waited_s - waited); diff --git a/zig/src/families/nemotron/cuda_reuse.zig b/zig/src/families/nemotron/cuda_reuse.zig index e73851a23c..180fa040a4 100644 --- a/zig/src/families/nemotron/cuda_reuse.zig +++ b/zig/src/families/nemotron/cuda_reuse.zig @@ -9,10 +9,11 @@ const Engine = @import("cuda_engine.zig").Engine; const Head = @import("cuda_mtp.zig").Head; const marks = @import("cuda_marks.zig"); -/// One kept state: its Mamba states in one device buffer, its key and value rows in another, where the engine stood. +/// A host state keeps both in `host` (key and value rows first); `mamba` is then an armed row's staging until `finish`. pub const Copy = struct { mamba: cuda.DeviceBuffer, kv: cuda.DeviceBuffer, + host: ?cuda.HostBuffer = null, // pinned: rows, then Mamba states; `mamba` is then an armed row's staging until `finish` head: bool, pos: usize, parity: usize, @@ -24,6 +25,11 @@ pub const Copy = struct { pub const Target = struct { e: *Engine, head: ?*Head, + + /// Kept states in pinned host memory (Engine.snap_host, a discrete card): the budget is the host's. + fn host(t: *const Target) bool { + return t.e.snap_host; + } }; /// `len` bytes at `live` in the bound sequence's planes; `mamba`: the SSM or conv state, which an armed chunk writes. @@ -94,9 +100,27 @@ fn move(t: *const Target, at: usize, copy: *const Copy, into: bool, mamba: bool) return off[0] + off[1]; } -/// Bytes a save at `at` takes, as the budget counts them: what `save` allocates, spare or new. +/// Bound sequence <-> host copy (`into`: to host), rows first; Mamba states too if `mamba`, else from `staging`. +fn hostMove(t: *const Target, at: usize, host: []u8, into: bool, mamba: bool, staging: ?u64) !void { + const o = t.e.ops(); + var store: [max_parts]Part = undefined; + var off = [2]usize{ 0, kvBytes(t.e.c, at, t.head != null) }; + for (parts(t.e.c, t.e.max_len, planes(t), at, &store)) |p| { + const at_host = off[@intFromBool(p.mamba)]; + const buf = host[at_host..][0..p.len]; + if (!p.mamba) { + if (into) try o.download(buf, p.live) else try o.upload(p.live, buf); + } else if (mamba) { + if (into) try o.download(buf, p.live) else try o.upload(p.live, buf); + } else if (staging) |st| try o.download(buf, st + (at_host - kvBytes(t.e.c, at, t.head != null))); // SSM, then conv + off[@intFromBool(p.mamba)] += p.len; + } +} + +/// Bytes a save at `at` takes, as the budget counts them: what `save` allocates, spare or new (host: its bytes). pub fn bytes(ptr: *anyopaque, at: u32) u64 { const t: *Target = @ptrCast(@alignCast(ptr)); + if (t.host()) return stateBytes(t.e.c, at, t.head != null); return mambaBytes(t.e.c) + allocBytes(kvBytes(t.e.c, at, t.head != null)); } @@ -104,9 +128,10 @@ pub fn bytes(ptr: *anyopaque, at: u32) u64 { pub fn ready(t: *Target, budget: u64) !u64 { const e = t.e; const len = mambaBytes(e.c); - e.spares.want = @min(marks.spares_max, budget / 4 / len); + // host states: one device staging buffer for an armed row, out of the device's room rather than the host budget + e.spares.want = if (t.host()) @min(marks.spares_max, 1) else @min(marks.spares_max, budget / 4 / len); e.spares.fill(e.ctx.d, len); - return e.spares.want * len; + return if (t.host()) 0 else e.spares.want * len; } /// Copy the live state after `at` prompt tokens. The pass is standing there. @@ -119,6 +144,15 @@ pub fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { if (e.armed.open) |span| if (at > span[0] and at < span[1]) return arm(t, at); if (e.pos != at) return error.NotAtMark; if (t.head) |h| if (h.pos != at) return error.NotAtMark; + if (t.host()) { + const copy = try e.gpa.create(Copy); + errdefer e.gpa.destroy(copy); + copy.* = .{ .mamba = .{ .ptr = 0, .len = 0, .d = e.ctx.d }, .kv = .{ .ptr = 0, .len = 0, .d = e.ctx.d }, .head = t.head != null, .pos = at, .parity = e.parity, .rows = e.rows }; + copy.host = try cuda.HostBuffer.alloc(e.ctx.d, stateBytes(e.c, at, copy.head)); + errdefer copy.host.?.free(); + try hostMove(t, at, copy.host.?.bytes, true, true, null); + return copy; + } const copy = try alloc(t, at, true); errdefer drop(ptr, copy); if (try move(t, at, copy, true, true) != stateBytes(e.c, at, t.head != null)) return error.SnapshotSize; @@ -132,7 +166,7 @@ fn alloc(t: *const Target, at: usize, kv: bool) !*Copy { errdefer e.gpa.destroy(copy); var m = e.spares.take() orelse try cuda.DeviceBuffer.alloc(e.ctx.d, mambaBytes(e.c)); errdefer if (!e.spares.give(m)) m.free(); - const rows = try cuda.DeviceBuffer.alloc(e.ctx.d, if (kv) allocBytes(kvBytes(e.c, at, t.head != null)) else 0); + const rows = try cuda.DeviceBuffer.alloc(e.ctx.d, if (kv and !t.host()) allocBytes(kvBytes(e.c, at, t.head != null)) else 0); copy.* = .{ .mamba = m, .kv = rows, .head = t.head != null, .pos = at, .parity = e.parity, .rows = e.rows }; return copy; } @@ -153,6 +187,19 @@ pub fn finish(t: Target) !void { const e = t.e; for (e.armed.marks()) |m| { const copy: *Copy = @ptrCast(@alignCast(m.copy)); + if (t.host()) { + // the chunk wrote this row's Mamba states to staging: rows and states to the host, then staging back + var h = cuda.HostBuffer.alloc(e.ctx.d, stateBytes(e.c, m.at, copy.head)) catch continue; // stays not whole + hostMove(&t, m.at, h.bytes, true, false, copy.mamba.ptr) catch { + h.free(); + continue; + }; + copy.host = h; + if (!e.spares.give(copy.mamba)) copy.mamba.free(); + copy.mamba = .{ .ptr = 0, .len = 0, .d = e.ctx.d }; + copy.whole = true; + continue; + } copy.kv = cuda.DeviceBuffer.alloc(e.ctx.d, allocBytes(kvBytes(e.c, m.at, copy.head))) catch continue; // stays not whole _ = try move(&t, m.at, copy, true, false); copy.whole = true; @@ -163,12 +210,17 @@ pub fn finish(t: Target) !void { /// What a kept state holds on the device, both buffers. pub fn held(saved: *anyopaque) usize { const copy: *Copy = @ptrCast(@alignCast(saved)); - return copy.mamba.len + copy.kv.len; + return copy.mamba.len + copy.kv.len + if (copy.host) |h| h.bytes.len else 0; } /// A kept state's bytes in `out` (its key and value rows, then its Mamba states), after the engine's stream. pub fn download(t: *const Target, saved: *anyopaque, out: []u8) !void { const copy: *Copy = @ptrCast(@alignCast(saved)); + if (copy.host) |h| { + try t.e.stream.synchronize(); + @memcpy(out[0..h.bytes.len], h.bytes); + return; + } const kv = kvBytes(t.e.c, copy.pos, copy.head); try t.e.ops().download(out[0..kv], copy.kv.ptr); try t.e.ops().download(out[kv..], copy.mamba.ptr); @@ -183,7 +235,10 @@ pub fn restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror! const e = t.e; if (copy.head != (t.head != null)) return error.NoHead; // checked before any copy: nothing half restored if (!copy.whole) return error.Incomplete; // its chunk failed before `finish` - _ = try move(t, copy.pos, copy, false, true); + if (copy.host) |h| { + try hostMove(t, copy.pos, h.bytes, false, true, null); + try e.stream.synchronize(); // the store may drop the host copy once this returns + } else _ = try move(t, copy.pos, copy, false, true); e.pos = copy.pos; e.parity = copy.parity; e.prev_keep = 0; @@ -196,8 +251,13 @@ pub fn drop(ptr: *anyopaque, saved: *anyopaque) void { const t: *Target = @ptrCast(@alignCast(ptr)); const copy: *Copy = @ptrCast(@alignCast(saved)); t.e.armed.forget(saved); // a copy evicted before its chunk ran: the chunk must not write it - if (!t.e.spares.give(copy.mamba)) copy.mamba.free(); - copy.kv.free(); + if (copy.host) |*h| { + // queued copies may still read or write it: they finish before the host memory goes + t.e.stream.synchronize() catch {}; + h.free(); + } + if (copy.mamba.len > 0 and !t.e.spares.give(copy.mamba)) copy.mamba.free(); + if (copy.kv.len > 0) copy.kv.free(); t.e.gpa.destroy(copy); } diff --git a/zig/src/families/nemotron/cuda_state.zig b/zig/src/families/nemotron/cuda_state.zig index 92bac057b9..f0bdb2626c 100644 --- a/zig/src/families/nemotron/cuda_state.zig +++ b/zig/src/families/nemotron/cuda_state.zig @@ -5,6 +5,7 @@ const lesson_route = @import("lesson_route.zig"); const cuda = @import("cuda"); const Config = @import("config.zig").Config; const kern = @import("cuda_kernels.zig"); +const glue = @import("cuda_glue.zig"); const sampler = @import("cuda_sampler.zig"); const seq_graphs = @import("cuda_seq_graphs.zig"); @@ -81,7 +82,7 @@ fn seqSizes(c: Config, max_len: usize) [seq_fields.len]usize { return .{ kv, kv, m[0], m[1], nm * 2 * W * cd * 2, nm * 2 * W * cd * 2, nm * 2 * W * c.mamba_heads * 4, 16, W * 4, W * c.hidden * 2, W * 4, @sizeOf(sampler.Rule) }; } -const scratch_count = 36; +const scratch_count = 39; /// Scratch windows and chunks share on one stream: window logits, chunk io, activations, the expert plan. fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { @@ -94,19 +95,21 @@ fn scratchSizes(c: Config, nch: usize) [scratch_count]usize { const xd: usize = c.inner(); const pairs = R * ns; const items = kern.maxItems(@intCast(pairs), c.experts + 2, 16); + const lat: usize = @max(c.latent, 1); // LatentMoE's rows in the latent width (Bolt), else unused return .{ - W * c.vocab * 2, 16, R * 4, - R * D * 2, c.vocab * 2, 16, - R * D * 2, R * D * 2, R * D * 2, - R * D * 2, R * D / 64 * 4, R * D * 2, - R * c.projDim() * 2, R * cd * 2, R * xd * 2, - R * xd * 2, R * xd / 64 * 4, R * c.qkvDim() * 2, - R * qd * 2, R * qd * 2, R * qd / 64 * 4, - W * nch * qd * 4, W * nch * c.heads * 4, W * nch * c.heads * 4, - 6 * R * c.experts * 4, R * ns * 4, R * ns * 4, - pairs * 4, items * 12, 8, - pairs * 4, (pairs + 1023) / 1024 * (c.experts + 2) * 4, pairs * c.expert_width * 2, - pairs * D * 4, shared_words * 4, R * kern.pocket_width * 4, + W * c.vocab * 2, 16, R * 4, + R * D * 2, c.vocab * 2, 16, + R * D * 2, R * D * 2, R * D * 2, + R * D * 2, R * D / 64 * 4, R * D * 2, + R * c.projDim() * 2, R * cd * 2, R * xd * 2, + R * xd * 2, R * xd / 64 * 4, R * c.qkvDim() * 2, + R * qd * 2, R * qd * 2, R * qd / 64 * 4, + W * nch * qd * 4, W * nch * c.heads * 4, W * nch * c.heads * 4, + glue.max_router_slices * R * c.experts * 4, R * ns * 4, R * ns * 4, + pairs * 4, items * 12, 8, + pairs * 4, (pairs + 1023) / 1024 * (c.experts + 2) * 4, pairs * c.expert_width * 2, + pairs * D * 4, shared_words * 4, R * kern.pocket_width * 4, + R * lat * 2, R * lat * 2, R * 4, }; } @@ -218,6 +221,9 @@ pub const Buffers = struct { rule: u64, // the sequence's sampled rule, as sample.cu reads it shared: u64, // the shared expert's fixed plans: 32 members, then (items, counts) for 1..16 rows pocket: u64, // the pocket's activations, fp32 [rows, pocket_width] + lat: u64, // LatentMoE: the rows fc1_latent_proj makes, bf16 [rows, latent] + z: u64, // LatentMoE: the routed experts' weighted sum, bf16 [rows, latent], fc2_latent_proj's input + zeros: u64, // LatentMoE: every pair's expert 0, the shared expert's picks as a one-expert layer /// Sizes every buffer for `max_len` cache rows and `nch` attention chunk partials a row; the own sequence's first. pub fn init(d: *const cuda.Driver, c: Config, max_len: usize, nch: usize, carve: ?*cuda.Carveout) !Buffers { @@ -262,7 +268,7 @@ pub const Buffers = struct { &b.logits, &b.p_meta, &b.p_ids, &b.p_hidden, &b.p_logits, &b.p_sampled, &b.emb, &b.h[0], &b.h[1], &b.y, &b.xs, &b.delta, &b.proj, &b.p_xc, &b.sy, &b.g, &b.gxs, &b.qkv, &b.q, &b.att, &b.axs, &b.po, &b.pm, &b.pl, &b.part, &b.pick, &b.wts, &b.plan.members, &b.plan.items, &b.plan.counts, - &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, + &b.plan.rank, &b.plan.hist, &b.act, &b.ymoe, &b.shared, &b.pocket, &b.lat, &b.z, &b.zeros, }; } @@ -285,6 +291,10 @@ pub const Buffers = struct { w[at..][0..8].* = .{ @intCast(c.experts), 0, n, @intCast(c.experts + 1), max_rows, n, 2, 2 }; } try ops.upload(b.shared, std.mem.sliceAsBytes(&w)); + if (c.latent > 0) { + const z: [pass_rows]u32 = @splat(0); + try ops.upload(b.zeros, std.mem.sliceAsBytes(&z)); + } try ops.s.synchronize(); } diff --git a/zig/src/families/nemotron/cuda_train_fwd.zig b/zig/src/families/nemotron/cuda_train_fwd.zig index 3b39c30976..e681d19845 100644 --- a/zig/src/families/nemotron/cuda_train_fwd.zig +++ b/zig/src/families/nemotron/cuda_train_fwd.zig @@ -86,7 +86,7 @@ pub fn experts(f: *const fwd.Forward, m: weights.MoE, rows: usize, scale: Scale, const o = f.ops; const ex = m.experts; const pairs = rows * c.slots(); - try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.routed_scaling, c.norm_topk); + try f.glue.route(b.y, m.router, m.bias, b.part, b.pick, b.wts, rows, c.hidden, c.experts, c.top_k, c.slots(), c.routed_scaling, c.norm_topk); const tile = cuda.experts.promptTile(ex.format); try o.plan(b.pick, pairs, ex.experts, tile, b.plan); const items = kern.maxItems(pairs, ex.experts, tile); diff --git a/zig/src/families/nemotron/cuda_vision.zig b/zig/src/families/nemotron/cuda_vision.zig new file mode 100644 index 0000000000..b017d534ad --- /dev/null +++ b/zig/src/families/nemotron/cuda_vision.zig @@ -0,0 +1,496 @@ +//! Bolt's image tower on CUDA: C-RADIO ViT-H/16 and the mlp1 projector, bf16 weights, fp32 residual; radio_vision.cu. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); + +const lt = cuda.cublaslt; +const DeviceBuffer = cuda.DeviceBuffer; +const lanes_ = @import("lanes"); +const radio = "vision_model.radio_model."; +const model = radio ++ "model."; + +pub const hidden = 1280; +pub const heads = 16; +pub const head_dim = 80; // radio_vision.cu's attention is built for 80 +pub const depth = 32; +pub const mlp = 5120; +pub const patch = 16; +pub const patch_width = 3 * patch * patch; // 768, (c, dy, dx) +pub const prefix_tokens = 10; // 3 cls + 7 registers, dropped after the last block +pub const cpe_side = 128; // pos_embed's 16384 rows +pub const shuffled = 4 * hidden; // pixel shuffle v2: 2x2 patches a row +pub const projector_hidden = 20480; +pub const out_hidden = 4096; // the language model's width +pub const ln_eps: f32 = 1e-6; +pub const rms_eps: f32 = 1e-5; +/// One call's patches at most: max_num_patches. +pub const max_patches = 13312; +/// The prompt's placeholder ids (tokenizer.json): `` repeated, between `` and ``. +pub const image_token: u32 = 18; +pub const image_start: u32 = 19; +pub const image_end: u32 = 20; +const workspace_bytes = 32 << 20; +/// The processor's normalization (vision/nemotron_image.zig's MEAN and STD): input_conditioner must hold the same. +const clip_mean = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +const clip_std = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; + +/// What config.json must say for this module to run the checkpoint's tower; anything else is refused. +pub fn check(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !void { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + // Bolt's quantization_config makes it ~6.5 MB + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 26)); + defer gpa.free(bytes); + return checkConfig(gpa, bytes); +} + +pub fn checkConfig(gpa: std.mem.Allocator, bytes: []const u8) !void { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, bytes, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const v = (root.get("vision_config") orelse return error.UnsupportedBoltVision).object; + const ints = [_]struct { std.json.ObjectMap, []const u8, i64 }{ + .{ v, "hidden_size", hidden }, .{ v, "num_attention_heads", heads }, .{ v, "num_hidden_layers", depth }, + .{ v, "patch_size", patch }, .{ v, "num_cls_tokens", 3 }, .{ v, "num_registers", 7 }, + .{ root, "projector_hidden_size", projector_hidden }, .{ root, "img_context_token_id", image_token }, .{ root, "vit_hidden_size", hidden }, + }; + for (ints) |w| if (num(w[0].get(w[1])) != w[2]) return error.UnsupportedBoltVision; + const llm = (root.get("llm_config") orelse return error.UnsupportedBoltVision).object; + if (num(llm.get("hidden_size")) != out_hidden) return error.UnsupportedBoltVision; + if (!isString(v.get("hidden_act"), "gelu") or !isString(root.get("ps_version"), "v2")) return error.UnsupportedBoltVision; + if (!isBool(v.get("use_swiglu_ffn"), false) or !isBool(root.get("use_thumbnail"), false) or !isBool(root.get("vision_final_layernorm"), false)) return error.UnsupportedBoltVision; + if (!isFloat(root.get("downsample_ratio"), 0.5) or !isFloat(v.get("layer_norm_eps"), 1e-6) or !isFloat(v.get("mlp_ratio"), 4.0)) return error.UnsupportedBoltVision; +} + +fn num(x: ?std.json.Value) i64 { + const v = x orelse return -1; + return if (v == .integer) v.integer else -1; +} + +fn isString(x: ?std.json.Value, want: []const u8) bool { + const v = x orelse return false; + return v == .string and std.mem.eql(u8, v.string, want); +} + +fn isBool(x: ?std.json.Value, want: bool) bool { + const v = x orelse return !want; // an absent flag is false + return v == .bool and v.bool == want; +} + +fn isFloat(x: ?std.json.Value, want: f64) bool { + const v = x orelse return false; + const f: f64 = switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => return false, + }; + return @abs(f - want) <= 1e-12 * @abs(want); +} + +const Block = struct { norm1_w: u64, norm1_b: u64, qkv_w: u64, qkv_b: u64, proj_w: u64, proj_b: u64, norm2_w: u64, norm2_b: u64, fc1_w: u64, fc1_b: u64, fc2_w: u64, fc2_b: u64 }; + +const Fns = struct { + embed: cuda.Function, + layernorm: cuda.Function, + qkv: cuda.Function, + attention: cuda.Function, + bias_gelu: cuda.Function, + bias_residual: cuda.Function, + shuffle_rms: cuda.Function, + relu2: cuda.Function, + cast: cuda.Function, +}; + +/// Where one call's stages live in the scratch (device pointers), for tests that read them back after `encode`. +pub const Stages = struct { + /// fp32 [prefix_tokens + n, 1280]: the residual stream, after `encode` the last block's output (tower_out). + x: u64, +}; + +/// The tower's and projector's weights and one call's scratch in one allocation, its kernels and GEMM plans. +pub const Tower = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + mod: cuda.Module, + f: Fns, + lib: lt.Library, + arena: DeviceBuffer, // weights 1.681 GB, scratch for max_patches 0.702 GB + at: usize = 0, + scratch: u64 = 0, + embed_w: u64 = 0, + pos: u64 = 0, + cls: u64 = 0, + blocks: [depth]Block = undefined, + norm_w: u64 = 0, + proj1_w: u64 = 0, + proj2_w: u64 = 0, + gemms: Gemms = undefined, + /// Blocks `encode` runs (tests: 0 leaves the embedded rows in Stages.x); the projector runs only after all 32. + run_depth: usize = depth, + + /// Loads vision_model. and mlp1. (bf16) from Bolt's `dir`, synchronously, into an arena sized for `max_patches`. + pub fn load(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8) !*Tower { + try check(gpa, io, dir); + const t = try gpa.create(Tower); + errdefer gpa.destroy(t); + t.* = .{ .gpa = gpa, .d = d, .mod = undefined, .f = undefined, .lib = undefined, .arena = undefined }; + t.lib = try lt.Library.open(); + errdefer t.lib.close(); + t.gemms = .{ .lib = &t.lib }; + errdefer t.gemms.deinit(); + t.mod = try cuda.Module.load(d, cuda.kernels.radio_vision); + errdefer t.mod.unload(); + inline for (@typeInfo(Fns).@"struct".field_names) |name| @field(t.f, name) = try t.mod.function("tf_rv_" ++ name); + t.arena = try DeviceBuffer.alloc(d, weightBytes() + scratchBytes(max_patches)); + errdefer t.arena.free(); + try t.arena.fill8(0, null); // touched once now: the driver commits memory on first touch, slowly + var vk = try core.Checkpoint.openModelPrefix(gpa, io, dir, "vision_model."); + defer vk.close(); + var pk = try core.Checkpoint.openModelPrefix(gpa, io, dir, "mlp1."); + defer pk.close(); + { // the input conditioner is the processor's normalization, already applied on the host: it must agree + const mean = try vk.expect(radio ++ "input_conditioner.norm_mean", .f32, &.{ 3, 1, 1 }); + const stdv = try vk.expect(radio ++ "input_conditioner.norm_std", .f32, &.{ 3, 1, 1 }); + for (0..3) |c| { + const m: f32 = @bitCast(std.mem.readInt(u32, mean.bytes[c * 4 ..][0..4], .little)); + const s: f32 = @bitCast(std.mem.readInt(u32, stdv.bytes[c * 4 ..][0..4], .little)); + if (m != clip_mean[c] or s != clip_std[c]) return error.UnsupportedBoltVision; + } + } + _ = try vk.expect("vision_model.summary_idxs", .i64, &.{2}); // the summary heads' cls rows: not used for images + _ = try vk.expect(model ++ "patch_generator.video_embedder.weight", .bf16, &.{ hidden, 2 * patch_width }); // video: checked, not loaded + t.embed_w = try t.get(&vk, model ++ "patch_generator.embedder.weight", &.{ hidden, patch_width }); + t.pos = try t.get(&vk, model ++ "patch_generator.pos_embed", &.{ 1, cpe_side * cpe_side, hidden }); + t.cls = try t.get(&vk, model ++ "patch_generator.cls_token.token", &.{ prefix_tokens, hidden }); + for (&t.blocks, 0..) |*b, i| { + const shapes = .{ + .{ "norm1_w", "norm1.weight", .{hidden} }, .{ "norm1_b", "norm1.bias", .{hidden} }, + .{ "qkv_w", "attn.qkv.weight", .{ 3 * hidden, hidden } }, .{ "qkv_b", "attn.qkv.bias", .{3 * hidden} }, + .{ "proj_w", "attn.proj.weight", .{ hidden, hidden } }, .{ "proj_b", "attn.proj.bias", .{hidden} }, + .{ "norm2_w", "norm2.weight", .{hidden} }, .{ "norm2_b", "norm2.bias", .{hidden} }, + .{ "fc1_w", "mlp.fc1.weight", .{ mlp, hidden } }, .{ "fc1_b", "mlp.fc1.bias", .{mlp} }, + .{ "fc2_w", "mlp.fc2.weight", .{ hidden, mlp } }, .{ "fc2_b", "mlp.fc2.bias", .{hidden} }, + }; + inline for (shapes) |sh| { + var nb: [128]u8 = undefined; + const shape: [sh[2].len]usize = sh[2]; + @field(b, sh[0]) = try t.get(&vk, try std.fmt.bufPrint(&nb, model ++ "blocks.{d}.{s}", .{ i, sh[1] }), &shape); + } + } + t.norm_w = try t.get(&pk, "mlp1.0.weight", &.{shuffled}); + t.proj1_w = try t.get(&pk, "mlp1.1.weight", &.{ projector_hidden, shuffled }); + t.proj2_w = try t.get(&pk, "mlp1.3.weight", &.{ out_hidden, projector_hidden }); + if (vk.unused() != 0 or pk.unused() != 0) return error.UnusedVisionTensors; + t.scratch = try t.take(scratchBytes(max_patches)); + return t; + } + + pub fn deinit(t: *Tower) void { + t.gemms.deinit(); + t.arena.free(); + t.mod.unload(); + t.lib.close(); + t.gpa.destroy(t); + } + + /// The weights' bytes in the arena (each allocation 256-byte aligned): 1,680,788,480. + pub fn weightBytes() usize { + const a = aligned; + var n: usize = a(hidden * patch_width * 2) + a(cpe_side * cpe_side * hidden * 2) + a(prefix_tokens * hidden * 2); + n += depth * (4 * a(hidden * 2) + a(3 * hidden * hidden * 2) + a(3 * hidden * 2) + a(hidden * hidden * 2) + a(hidden * 2) + a(mlp * hidden * 2) + a(mlp * 2) + a(hidden * mlp * 2) + a(hidden * 2)); + return n + a(shuffled * 2) + a(projector_hidden * shuffled * 2) + a(out_hidden * projector_hidden * 2); + } + + /// The scratch one call of `n` patches needs, in bytes (aligned pieces; 13,312 patches: 702 MB). + pub fn scratchBytes(n: usize) usize { + const a = aligned; + const rows = prefix_tokens + n; + const m4 = n / 4; + return a(n * patch_width * 2) + // bf16 patches + a(rows * hidden * 4) + // x, the fp32 residual + a(@max(rows * hidden, m4 * shuffled) * 2) + // h: LayerNorm out, then the shuffled rows + a(@max(rows * @max(3 * hidden, mlp), m4 * projector_hidden) * 4) + // fp32 GEMM out + 3 * a(rows * hidden * 2) + // q, k, v + a(rows * hidden * 2) + // attention out + a(@max(rows * mlp, m4 * projector_hidden) * 2) + // GELU / relu^2 rows + a(n * 4 * 4) * 2 + // CPE taps and weights + workspace_bytes; + } + + fn take(t: *Tower, len: usize) !u64 { + const at = aligned(t.at); + if (at + len > t.arena.len) return error.ArenaExhausted; + t.at = at + len; + return t.arena.ptr + at; + } + + fn get(t: *Tower, ck: *core.Checkpoint, name: []const u8, shape: []const usize) !u64 { + const tn = try ck.expect(name, .bf16, shape); + const p = try t.take(tn.bytes.len); + try t.arena.upload(p - t.arena.ptr, tn.bytes); + return p; + } + + /// One image's patches (fp32 [gh * gw, 768]) to bf16 [gh * gw / 4, 4096] rows at `out` on `s`; inputs copied. + pub fn encode(t: *Tower, s: cuda.Stream, patches: []const f32, gh: u32, gw: u32, out: u64) !Stages { + const n = @as(usize, gh) * gw; + if (gh == 0 or gw == 0 or gh % 2 != 0 or gw % 2 != 0 or n > max_patches) return error.BadImageGrid; + if (patches.len != n * patch_width) return error.BadImagePatches; + const gpa = t.gpa; + const rows = prefix_tokens + n; + const m4 = n / 4; + const px16 = try gpa.alloc(u16, n * patch_width); + defer gpa.free(px16); + for (px16, patches) |*o, v| o.* = bf16(v); + const idx = try gpa.alloc(i32, n * 4); + defer gpa.free(idx); + const wts = try gpa.alloc(f32, n * 4); + defer gpa.free(wts); + cpeTaps(gh, gw, idx, wts); + + var cur = t.scratch; + const end = t.scratch + scratchBytes(max_patches); + const d_px = carve(&cur, n * patch_width * 2); + const x = carve(&cur, rows * hidden * 4); + const h = carve(&cur, @max(rows * hidden, m4 * shuffled) * 2); + const acc = carve(&cur, @max(rows * @max(3 * hidden, mlp), m4 * projector_hidden) * 4); + const q = carve(&cur, rows * hidden * 2); + const k = carve(&cur, rows * hidden * 2); + const v = carve(&cur, rows * hidden * 2); + const att = carve(&cur, rows * hidden * 2); + const mid = carve(&cur, @max(rows * mlp, m4 * projector_hidden) * 2); + const d_idx = carve(&cur, n * 4 * 4); + const d_wts = carve(&cur, n * 4 * 4); + const ws = cur; + const ws_len = end - cur; + try t.upload(s, d_px, std.mem.sliceAsBytes(px16)); + try t.upload(s, d_idx, std.mem.sliceAsBytes(idx)); + try t.upload(s, d_wts, std.mem.sliceAsBytes(wts)); + + const mm = &t.gemms; + try mm.run(n, hidden, patch_width, d_px, t.embed_w, acc, ws, ws_len, s); + try go(t.f.embed, s, rows, 1, 256, .{ acc, t.pos, d_idx, d_wts, t.cls, x, int(prefix_tokens), int(hidden) }); + const scale_log2: f32 = std.math.log2e / @sqrt(@as(f32, head_dim)); + for (t.blocks[0..t.run_depth]) |b| { + try go(t.f.layernorm, s, rows, 1, 256, .{ x, b.norm1_w, b.norm1_b, h, int(hidden), ln_eps }); + try mm.run(rows, 3 * hidden, hidden, h, b.qkv_w, acc, ws, ws_len, s); + try go(t.f.qkv, s, rows, 1, 256, .{ acc, b.qkv_b, q, k, v, int(rows), int(heads), int(head_dim) }); + try go(t.f.attention, s, (rows + 63) / 64, heads, 128, .{ q, k, v, att, int(rows), int(heads), scale_log2 }); + try mm.run(rows, hidden, hidden, att, b.proj_w, acc, ws, ws_len, s); + try go(t.f.bias_residual, s, rows, 1, 256, .{ acc, b.proj_b, x, int(hidden) }); + try go(t.f.layernorm, s, rows, 1, 256, .{ x, b.norm2_w, b.norm2_b, h, int(hidden), ln_eps }); + try mm.run(rows, mlp, hidden, h, b.fc1_w, acc, ws, ws_len, s); + try go(t.f.bias_gelu, s, rows, 1, 256, .{ acc, b.fc1_b, mid, int(mlp) }); + try mm.run(rows, hidden, mlp, mid, b.fc2_w, acc, ws, ws_len, s); + try go(t.f.bias_residual, s, rows, 1, 256, .{ acc, b.fc2_b, x, int(hidden) }); + } + if (t.run_depth < depth) { + try s.synchronize(); + return .{ .x = x }; + } + // projector: the patch rows (prefix dropped) shuffled 2x2, RMSNorm, Linear, relu^2, Linear + try go(t.f.shuffle_rms, s, m4, 1, 256, .{ x, t.norm_w, h, int(prefix_tokens), int(gw), int(hidden), rms_eps }); + try mm.run(m4, projector_hidden, shuffled, h, t.proj1_w, acc, ws, ws_len, s); + try go(t.f.relu2, s, m4, 1, 256, .{ acc, mid, int(projector_hidden) }); + try mm.run(m4, out_hidden, projector_hidden, mid, t.proj2_w, acc, ws, ws_len, s); + try go(t.f.cast, s, m4, 1, 256, .{ acc, out, int(out_hidden) }); + try s.synchronize(); // the host's patch rows and tables are freed on return + return .{ .x = x }; + } + + fn upload(t: *Tower, s: cuda.Stream, dst: u64, bytes: []const u8) !void { + if (bytes.len == 0) return; + try t.d.check(t.d.api.cuMemcpyHtoDAsync_v2(dst, bytes.ptr, bytes.len, s.handle), "cuMemcpyHtoDAsync"); + } +}; + +fn go(f: cuda.Function, s: cuda.Stream, gx: usize, gy: usize, block: u32, args: anytype) !void { + var a: cuda.Args = .{}; + inline for (args) |v| a.add(v); + try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(gx), .y = @intCast(gy) }, .block = .{ .x = block } }, s, &a); +} + +fn carve(p: *u64, len: usize) u64 { + const r = p.*; + p.* += aligned(len); + return r; +} + +/// cuBLASLt plans for the shapes a call runs, made on first use and kept (a new image size adds plans). +const Gemms = struct { + lib: *const lt.Library, + plans: [48]?struct { m: usize, n: usize, k: usize, lin: lt.Linear } = @splat(null), + next: usize = 0, + + fn run(g: *Gemms, m: usize, n: usize, k: usize, x: u64, w: u64, out: u64, ws: u64, ws_len: usize, s: cuda.Stream) !void { + for (&g.plans) |*p| if (p.*) |*have| { + if (have.m == m and have.n == n and have.k == k) return have.lin.run(x, w, out, ws, ws_len, s.handle); + }; + // a full table drops its oldest plan (8 shapes an image size: six images' worth kept) + const slot = &g.plans[g.next]; + if (slot.*) |*old| old.lin.deinit(); + slot.* = .{ .m = m, .n = n, .k = k, .lin = try lt.Linear.init(g.lib, m, n, k, .f32, ws_len) }; + g.next = (g.next + 1) % g.plans.len; + return slot.*.?.lin.run(x, w, out, ws, ws_len, s.handle); + } + + fn deinit(g: *Gemms) void { + for (&g.plans) |*p| if (p.*) |*have| { + have.lin.deinit(); + p.* = null; + }; + } +}; + +fn aligned(n: usize) usize { + return std.mem.alignForward(usize, n, 256); +} + +fn int(x: usize) c_int { + return @intCast(x); +} + +/// fp32 to bf16, round to nearest even (torch's .to(bfloat16)). +fn bf16(v: f32) u16 { + const b: u32 = @bitCast(v); + if (std.math.isNan(v)) return 0x7fc0; + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// CPE taps: each patch's 4 rows and weights in the 128x128 table resized to max(gh, gw) and cropped, as aten (fp32). +pub fn cpeTaps(gh: u32, gw: u32, idx: []i32, wts: []f32) void { + const m = @max(gh, gw); + for (0..gh) |i| { + const hy = axis(@intCast(i), m); + for (0..gw) |j| { + const wx = axis(@intCast(j), m); + const p = i * gw + j; + idx[p * 4 + 0] = hy.lo * cpe_side + wx.lo; + idx[p * 4 + 1] = hy.lo * cpe_side + wx.hi; + idx[p * 4 + 2] = hy.hi * cpe_side + wx.lo; + idx[p * 4 + 3] = hy.hi * cpe_side + wx.hi; + wts[p * 4 + 0] = wx.wlo; + wts[p * 4 + 1] = wx.whi; + wts[p * 4 + 2] = hy.wlo; + wts[p * 4 + 3] = hy.whi; + } + } +} + +const Axis = struct { lo: i32, hi: i32, wlo: f32, whi: f32 }; + +/// aten's area_pixel_compute_source_index (linear, align_corners=False) and lambdas for output `o` of `size`. +fn axis(o: u32, size: u32) Axis { + const scale: f32 = @as(f32, cpe_side) / @as(f32, @floatFromInt(size)); + const src = @max(scale * (@as(f32, @floatFromInt(o)) + 0.5) - 0.5, 0.0); + const lo: i32 = @intFromFloat(src); + const hi: i32 = if (lo < cpe_side - 1) lo + 1 else lo; + const whi = @min(@max(src - @as(f32, @floatFromInt(lo)), 0.0), 1.0); + return .{ .lo = lo, .hi = hi, .wlo = 1.0 - whi, .whi = whi }; +} + +test "CPE taps: a 128 grid reads the table as is; an upsampled one interpolates and clamps at the far edge" { + var idx: [4 * 4]i32 = undefined; + var wts: [4 * 4]f32 = undefined; + const a = axis(5, cpe_side); + try std.testing.expectEqual(@as(i32, 5), a.lo); + try std.testing.expectEqual(@as(f32, 0), a.whi); + const last = axis(255, 256); + try std.testing.expectEqual(@as(i32, 127), last.lo); + try std.testing.expectEqual(@as(i32, 127), last.hi); + cpeTaps(2, 2, &idx, &wts); + // m = 2: scale 64, source 31.5 for row/col 0 + try std.testing.expectEqual(@as(i32, 31 * cpe_side + 31), idx[0]); + try std.testing.expectEqual(@as(f32, 0.5), wts[0]); +} + +test "Bolt's config is accepted; another tower or projector is refused" { + const gpa = std.testing.allocator; + const good = + \\{"vision_config":{"hidden_size":1280,"num_attention_heads":16,"num_hidden_layers":32,"patch_size":16, + \\"num_cls_tokens":3,"num_registers":7,"hidden_act":"gelu","use_swiglu_ffn":false,"layer_norm_eps":1e-06, + \\"mlp_ratio":4.0},"projector_hidden_size":20480,"img_context_token_id":18,"vit_hidden_size":1280, + \\"llm_config":{"hidden_size":4096},"ps_version":"v2","use_thumbnail":false,"vision_final_layernorm":false, + \\"downsample_ratio":0.5} + ; + try checkConfig(gpa, good); + const cases = [_][2][]const u8{ .{ "\"num_registers\":7", "\"num_registers\":4" }, .{ "\"ps_version\":\"v2\"", "\"ps_version\":\"v1\"" }, .{ "\"use_swiglu_ffn\":false", "\"use_swiglu_ffn\":true" }, .{ "\"downsample_ratio\":0.5", "\"downsample_ratio\":0.25" } }; + for (cases) |c| { + const bad = try std.mem.replaceOwned(u8, gpa, good, c[0], c[1]); + defer gpa.free(bad); + try std.testing.expectError(error.UnsupportedBoltVision, checkConfig(gpa, bad)); + } +} + +/// An image prompt's rows for the prompt pass: its images' features, and each prompt row's feature row (-1: text). +pub const Media = struct { + gpa: std.mem.Allocator, + features: DeviceBuffer, // bf16 [image rows, out_hidden] + src: DeviceBuffer, // i32 [prompt] + + /// `images` (the request's, in prompt order) through the tower on `s`; each fills its ... run. + pub fn init(gpa: std.mem.Allocator, t: *Tower, s: cuda.Stream, prompt: []const u32, images: []const lanes_.Image) !Media { + const host = try gpa.alloc(i32, prompt.len); + defer gpa.free(host); + const rows = try placeholders(prompt, images, host); + var m: Media = .{ .gpa = gpa, .features = try DeviceBuffer.alloc(t.d, @max(rows, 1) * out_hidden * 2), .src = undefined }; + errdefer m.features.free(); + m.src = try DeviceBuffer.alloc(t.d, prompt.len * 4); + errdefer m.src.free(); + var at: usize = 0; + for (images) |im| { + const gh: u32 = @intCast(im.grid[1]); + const gw: u32 = @intCast(im.grid[2]); + _ = try t.encode(s, im.patches, gh, gw, m.features.ptr + at * out_hidden * 2); + at += @as(usize, gh) * gw / 4; + } + try m.src.upload(0, std.mem.sliceAsBytes(host)); + return m; + } + + pub fn deinit(m: *Media) void { + m.features.free(); + m.src.free(); + } +}; + +/// Each image's image_token run between image_start and image_end; `src` gets each row's feature row (-1: text). +pub fn placeholders(prompt: []const u32, images: []const lanes_.Image, src: []i32) !usize { + @memset(src, -1); + var next: usize = 0; // feature rows handed out + var k: usize = 0; // images placed + var i: usize = 0; + while (i < prompt.len) : (i += 1) { + if (prompt[i] != image_token) continue; + if (k == images.len or i == 0 or prompt[i - 1] != image_start) return error.BadImagePlaceholders; + const im = images[k]; + if (im.grid[0] != 1 or im.grid[1] <= 0 or im.grid[2] <= 0 or @mod(im.grid[1] * im.grid[2], 4) != 0) return error.BadImageGrid; + const n: usize = @intCast(@divExact(im.grid[1] * im.grid[2], 4)); + if (im.patches.len != n * 4 * patch_width) return error.BadImagePatches; + if (i + n >= prompt.len or prompt[i + n] != image_end) return error.BadImagePlaceholders; + for (prompt[i..][0..n], src[i..][0..n], 0..) |id, *x, j| { + if (id != image_token) return error.BadImagePlaceholders; + x.* = @intCast(next + j); + } + next += n; + k += 1; + i += n; + } + if (k != images.len) return error.BadImagePlaceholders; + return next; +} + +test "placeholder runs map to feature rows in order; a short run is refused" { + const px: [8 * patch_width]f32 = @splat(0); + const ims = [_]lanes_.Image{ .{ .patches = &px, .grid = .{ 1, 2, 2 } }, .{ .patches = &px, .grid = .{ 1, 2, 4 } } }; + const p = [_]u32{ 1, image_start, image_token, image_end, 7, image_start, image_token, image_token, image_end, 9 }; + var src: [p.len]i32 = undefined; + // the first image is 4 patches (1 token) but its patches slice is 8: refused; give each its own size + const one: [4 * patch_width]f32 = @splat(0); + const good = [_]lanes_.Image{ .{ .patches = &one, .grid = .{ 1, 2, 2 } }, .{ .patches = &px, .grid = .{ 1, 2, 4 } } }; + try std.testing.expectEqual(@as(usize, 3), try placeholders(&p, &good, &src)); + try std.testing.expectEqualSlices(i32, &.{ -1, -1, 0, -1, -1, -1, 1, 2, -1, -1 }, &src); + try std.testing.expectError(error.BadImagePatches, placeholders(&p, &ims, &src)); + const short = [_]u32{ image_start, image_token, image_end, image_start, image_token, image_end }; + try std.testing.expectError(error.BadImagePlaceholders, placeholders(&short, &good, &src)); +} diff --git a/zig/src/families/nemotron/cuda_weights.zig b/zig/src/families/nemotron/cuda_weights.zig index ff25a8da37..26a3ba6834 100644 --- a/zig/src/families/nemotron/cuda_weights.zig +++ b/zig/src/families/nemotron/cuda_weights.zig @@ -28,7 +28,10 @@ pub const Attention = struct { qkv: QLinear, o: QLinear, o_rest: u64 = 0, adapte /// The pocket kernels' neurons (config's pocket): input weights and outputs, [n, D] fp32 each, from the 4-bit codes. pub const Pocket = kern.Pocket; /// The shared expert's down projection rest in two halves (experts E and E + 1); its learned change reads both. -pub const MoE = struct { router: u64, bias: u64, experts: Experts, rest: [2]u64 = .{ 0, 0 }, adapter: ?Adapter = null, pocket: ?Pocket = null }; +pub const MoE = struct { router: u64, bias: u64, experts: Experts, rest: [2]u64 = .{ 0, 0 }, adapter: ?Adapter = null, pocket: ?Pocket = null, latent: ?Latent = null }; + +/// LatentMoE's parts (Bolt): latent-width projections around the routed `experts`; `shared` reads full-width rows. +pub const Latent = struct { fc1: QLinear, fc2: QLinear, shared: Experts }; pub const Block = struct { kind: Kind, norm: u64, mamba: Mamba = undefined, attn: Attention = undefined, moe: MoE = undefined }; pub const Mtp = struct { enorm: u64, hnorm: u64, eh_proj: QLinear, attn_norm: u64, attn: Attention, moe_norm: u64, moe: MoE, final_norm: u64 }; @@ -458,6 +461,12 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, var ck = try core.Checkpoint.openModel(gpa, io, dir); defer ck.close(); try src.add(&ck); + if (c.omni) { + // Bolt: the language model under language_model., its vision tower and projector beside it + try ck.skip("vision_model."); + try ck.skip("mlp1."); + ck.scope = "language_model."; + } w.blocks = try gpa.alloc(Block, c.layers); for (c.kinds[0..c.layers], 0..) |kind, i| { var nm: [128]u8 = undefined; @@ -470,7 +479,7 @@ pub fn load(gpa: std.mem.Allocator, io: std.Io, ops: kern.Ops, dir: []const u8, switch (kind) { .mamba => b.mamba = try L.mamba(&ck, try join(&nb, name, ".mamba"), mixer), .attention => b.attn = try L.attention(&ck, try join(&nb, name, ".attn"), mixer), - .moe => b.moe = try L.moe(&ck, try join(&nb, name, ".moe"), mixer, c), + .moe => b.moe = if (c.latent > 0) try modelopt.latent(&L, &ck, try join(&nb, name, ".moe"), mixer, c, i) else try L.moe(&ck, try join(&nb, name, ".moe"), mixer, c), } w.blocks[i] = b; } diff --git a/zig/src/families/nemotron/glue_ref.zig b/zig/src/families/nemotron/glue_ref.zig index f4fc4550b9..d201f98be9 100644 --- a/zig/src/families/nemotron/glue_ref.zig +++ b/zig/src/families/nemotron/glue_ref.zig @@ -159,8 +159,8 @@ pub fn router(x: []const u16, w: []const u16, part: []f32, rows: usize, d: usize /// Sum slices in order, rank sigmoid scores plus bias (lowest id on ties), then add shared slots. pub fn topk(part: []const f32, bias: []const f32, idx: []i32, wt: []f32, rows: usize, scaling: f32, e: usize, sk: usize, top_k: usize, ns: usize, norm: bool) void { - var score: [256]f32 = undefined; - var sel: [256]f32 = undefined; + var score: [512]f32 = undefined; // the top-k kernel's MAXC: up to 512 experts + var sel: [512]f32 = undefined; for (0..rows) |r| { for (0..e) |ex| { var logit: f32 = 0; @@ -170,8 +170,8 @@ pub fn topk(part: []const f32, bias: []const f32, idx: []i32, wt: []f32, rows: u sel[ex] = if (std.math.isNan(v)) -std.math.inf(f32) else v; } var total: f32 = 0; - var probs: [8]f32 = undefined; - var ids: [8]i32 = undefined; + var probs: [16]f32 = undefined; + var ids: [16]i32 = undefined; for (0..top_k) |k| { var bi: usize = 0; for (1..e) |ex| if (sel[ex] > sel[bi]) { diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..f77d0bedc2 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -1,4 +1,8 @@ //! CUDA family registration, device setup, memory admission and the native lane host. + +/// The serve flags this build's engine takes past the common ones (server/cli.zig served()): two ranks. +pub const serves: []const []const u8 = &.{"--tp"}; + const std = @import("std"); const cuda = @import("cuda"); const api = @import("engine_api"); @@ -8,6 +12,8 @@ const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); const nemotron_slide = @import("nemotron_slide.zig"); +const aliases = @import("cuda_alias.zig").aliases; +const reads = @import("cuda_alias.zig").reads; /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{nemotron.native}; @@ -22,8 +28,14 @@ const nemotron_snaps: api.prompt_cache.Snapshots.VTable = .{ pub const backends: []const []const u8 = &.{"cuda"}; 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 }; + var n: usize = 0; + for (registry) |F| n += 1 + aliases(F).len; + var out: [n]api.Family = undefined; + var i: usize = 0; + for (registry) |F| for ([_][]const u8{F.model_type} ++ aliases(F)) |t| { + out[i] = .{ .model_type = t, .formats = F.formats }; + i += 1; + }; const final = out; break :blk &final; }; @@ -67,10 +79,11 @@ fn setting(a: Allocator, flag: ?u32, name: [:0]const u8, default: u32, problem: /// The model's window (config.json's max_position_embeddings, text_config's first), 0 when it names none. fn modelContext(a: Allocator, io: std.Io, dir: []const u8) i64 { const path = std.fs.path.join(a, &.{ dir, "config.json" }) catch return 0; - const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(16 << 20)) catch return 0; + const bytes = std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(32 << 20)) catch return 0; const doc = std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}) catch return 0; if (doc != .object) return 0; - const text = if (doc.object.get("text_config")) |t| (if (t == .object) t else doc) else doc; + const inner = doc.object.get("text_config") orelse doc.object.get("llm_config"); // Bolt's is llm_config + const text = if (inner) |t| (if (t == .object) t else doc) else doc; const limit = text.object.get("max_position_embeddings") orelse doc.object.get("max_position_embeddings") orelse return 0; return if (limit == .integer and limit.integer > 0) limit.integer else 0; } @@ -152,6 +165,14 @@ const Host = struct { slide: ?*anyopaque = null, // that learner's adapter, freed by slide_free before the family slide_free: ?*const fn (Allocator, *anyopaque) void = null, dir: []u8 = &.{}, // the model folder the learner writes its lessons into + follow_fn: ?*const fn (*anyopaque) anyerror!void = null, // rank 1 of two: the family's replay loop + + /// Rank 1: the family replays rank 0's calls on this thread (the context was made current here at open). + fn followFamily(p: *anyopaque) anyerror!void { + const h: *Host = @ptrCast(@alignCast(p)); + if (h.gpu) |g| try g.ctx.makeCurrent(); + return h.follow_fn.?(h.family); + } fn close(p: *anyopaque) void { const h: *Host = @ptrCast(@alignCast(p)); @@ -241,6 +262,21 @@ const Host = struct { const refusals = @import("cuda_refusal.zig"); +/// A tensor-parallel rank above 0: the family opened as that rank, then its replay loop; false with ``problem`` set. +pub fn follow(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !bool { + const opened = (try open(a, gpa, io, o, problem)) orelse return false; + defer opened.close(opened.ctx); + const f = opened.follow orelse { + problem.* = "this family has no tensor-parallel ranks"; + return false; + }; + f(opened.ctx) catch |e| { + problem.* = try std.fmt.allocPrint(a, "rank {d} stopped: {s}", .{ o.rank, @errorName(e) }); + return false; + }; + return true; +} + /// The engine for `o.dir`, or null with `problem` set when `refusal` names a flag or no CUDA family reads it. pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]const u8) !?api.Opened { if (refusals.backend(o)) |why| { @@ -248,7 +284,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c return null; } inline for (registry) |F| { - if (std.mem.eql(u8, o.model_type, F.model_type)) return openWith(F, a, gpa, io, o, problem); + if (reads(F, o.model_type)) return openWith(F, a, gpa, io, o, problem); } problem.* = try std.fmt.allocPrint(a, "the native CUDA engine has no backend for {s} checkpoints yet", .{o.model_type}); return null; @@ -322,20 +358,42 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }, }; const before = try pool(a, io, &g.ctx, problem) orelse return null; - const weights = weightBytes(io, o.dir); + const two = @hasDecl(F, "two_ranks") and F.two_ranks; + if (two and o.tp != 2) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on two GPUs here: serve it with --tp 2 --rank 0 on one host and --tp 2 --rank 1 --master on the other", .{F.model_type}); + return null; + } + if (!two and o.tp != 1) { + problem.* = try std.fmt.allocPrint(a, "{s} runs on one GPU: drop --tp", .{F.model_type}); + return null; + } + const weights = if (@hasDecl(F, "deviceWeightBytes")) F.deviceWeightBytes(io, o.dir) else weightBytes(io, o.dir); const held0 = cuda.usage(false).device; 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; } - const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments, .slide = o.slide }) catch |e| { + var fo: F.Options = .{ .context = @intCast(window), .drafts = o.drafts, .segments = segments }; + if (@hasField(F.Options, "vision")) fo.vision = o.vision; + if (@hasField(F.Options, "kv_bits")) fo.kv_bits = o.kv_bits else if (o.kv_bits != 8) { + problem.* = try std.fmt.allocPrint(a, "{s} keeps an int8 attention cache only: drop --kv-dtype", .{F.model_type}); + return null; + } + if (@hasField(F.Options, "slide")) fo.slide = o.slide; + if (@hasField(F.Options, "rank")) { + fo.tp = @intCast(o.tp); + fo.rank = @intCast(o.rank); + fo.master = if (o.master.len > 0) o.master else null; + fo.master_port = o.master_port; + } + const loaded = F.open(gpa, io, &g.ctx, o.dir, kernels, fo) catch |e| { problem.* = try std.fmt.allocPrint(a, "the native CUDA engine cannot load {s} with {s} glue kernels ({s})", .{ o.dir, kernels orelse "its own", @errorName(e) }); return null; }; defer if (!opened) loaded.deinit(loaded.ctx); const model = cuda.usage(false).device - held0; const after = try pool(a, io, &g.ctx, problem) orelse return null; - const room = after.room(model) -| loaded.learn_bytes; + const room = after.room(model) -| learnBytes(loaded); const free = loaded.free_streams; const streams = budget.admit(room, loaded.stream_bytes, free, o.lanes, o.lanes_fixed) catch |e| { problem.* = switch (e) { @@ -344,16 +402,23 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O }; return null; }; + if (@hasDecl(F, "setStreams")) F.setStreams(loaded.ctx, streams); + if (@hasDecl(F, "setRoom")) F.setRoom(loaded.ctx, room) catch |e| { // a family that sizes each stream's caches out of one pool + problem.* = try std.fmt.allocPrint(a, "the ranks could not agree on the stream pool ({s})", .{@errorName(e)}); + return null; + }; const h = try gpa.create(Host); errdefer gpa.destroy(h); h.* = .{ .gpa = gpa, .gpu = g, .family = loaded.ctx, .release = loaded.deinit, .inner = loaded.backend, .vtable = undefined, .cfg = undefined, .clock = undefined, .core = undefined, .host = undefined, .lone = loaded.lone }; - const cache = budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); - h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in passes of up to {d} rows{s}{s}; {s}", .{ - capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (loaded.learn_bytes > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(loaded.learn_bytes)}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", - if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache", .{toGib(cache)}) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", + // the lane core's kept states (Nemotron's snapshots); a family that keeps its own prompt states has no `target` + const host_snaps = @hasField(@TypeOf(loaded), "host_snaps") and loaded.host_snaps; + const cache = if (!@hasField(@TypeOf(loaded), "target")) 0 else if (host_snaps) budget.hostCacheBytes(a, io, o.prompt_cache_gib) else budget.cacheBytes(o.prompt_cache_gib, room, loaded.stream_bytes * (streams -| free)); + h.startup = try std.fmt.allocPrint(gpa, "CUDA sm_{d} device {d} ({s}{s}): model {d:.2} GiB{s}; {d} stream{s} at once, {d:.2} GiB each at a {d}-token window, of {d:.1} GiB left after a {d:.1} GiB reserve; prompts in {d}-row chunks{s}{s}; {s}", .{ + capability, device, name, if (after.unified) ", memory shared with the host" else "", toGib(model), if (learnBytes(loaded) > 0) try std.fmt.allocPrint(a, " and {d:.2} GiB kept for --slide's trainer", .{toGib(learnBytes(loaded))}) else "", streams, if (streams == 1) "" else "s", toGib(loaded.stream_bytes), window, toGib(room), toGib(after.reserve), F.prompt_rows, if (segments > 1) try std.fmt.allocPrint(a, ", {d} staggered segments a call", .{segments}) else "", + if (cache > 0) try std.fmt.allocPrint(a, ", a {d:.1} GiB prompt cache{s}", .{ toGib(cache), if (host_snaps) " in host memory" else "" }) else ", no prompt cache", if (kernels) |dir| try std.fmt.allocPrint(a, "glue kernels captured at {s}", .{dir}) else "own glue kernels", }); errdefer gpa.free(h.startup); - if (cache > 0) { + if (@hasField(@TypeOf(loaded), "target") and cache > 0) { const pc = api.prompt_cache; const store = try gpa.create(pc.Store); // states keep at any row at no prompt call (the head reads one token on); replies prefill when it is quiet @@ -364,31 +429,43 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O store.deinit(); gpa.destroy(store); }; - if (loaded.learns != null) { - const Slide = nemotron_slide.Adapter(F.slide); - h.dir = try gpa.dupe(u8, o.dir); - const s = try gpa.create(Slide); - s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); - h.slide = s; - h.slide_free = struct { - fn f(al: Allocator, p: *anyopaque) void { - const x: *Slide = @ptrCast(@alignCast(p)); - x.deinit(); - al.destroy(x); - } - }.f; - h.learn = s.hook(); + if (comptime @hasField(@TypeOf(loaded), "learns")) { // a family with a Living Weights learner + if (loaded.learns != null) { + const Slide = nemotron_slide.Adapter(F.slide); + h.dir = try gpa.dupe(u8, o.dir); + const s = try gpa.create(Slide); + s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); + h.slide = s; + h.slide_free = struct { + fn f(al: Allocator, p: *anyopaque) void { + const x: *Slide = @ptrCast(@alignCast(p)); + x.deinit(); + al.destroy(x); + } + }.f; + h.learn = s.hook(); + } } errdefer if (h.slide) |s| { h.slide_free.?(gpa, s); gpa.free(h.dir); }; // the family cuts its own prompt grid from position 0, as `tensorfold run` does: prefill_step 0 - try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null }, .{ .ctx = loaded.ctx, .text = F.explain }); + try h.serve(io, loaded.facts, loaded.rows, .{ .lanes = streams, .context_window = @intCast(window), .startup = h.startup, .prompt_cache = h.store != null, .warm_turns = h.store != null, .logprobs = loaded.backend.vtable.first_row != null, .prompt_logprobs = loaded.backend.vtable.prompt_rows, .structures = loaded.backend.vtable.first_masked != null, .vision = o.vision and @hasField(F.Options, "vision"), .vision_glm = @hasDecl(F, "vision_glm") and F.vision_glm, .vision_nemotron = @hasDecl(F, "vision_nemotron") and F.vision_nemotron, .plain_only = loaded.rows == 1, .admit = if (@hasDecl(F, "fits")) .{ .ctx = loaded.ctx, .fits = F.fits } else null }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; + const follow_loop: ?*const fn (*anyopaque) anyerror!void = if (@hasField(@TypeOf(loaded), "follow")) loaded.follow else null; + if (follow_loop) |f| { + h.follow_fn = f; + return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h, .follow = Host.followFamily }; + } return .{ .engine = h.host.engine(), .close = Host.close, .ctx = h }; } +/// A family's bytes kept for --slide (0 for families without a learner). +fn learnBytes(loaded: anytype) u64 { + return if (@hasField(@TypeOf(loaded), "learn_bytes")) loaded.learn_bytes else 0; +} + fn toGib(bytes: u64) f64 { return @as(f64, @floatFromInt(bytes)) / gib; } @@ -414,9 +491,13 @@ test "chip classes name the compute capability as gate entries do" { try std.testing.expectEqualStrings("nvidia-sm121", name); } -test "every registered family is listed for capabilities" { - try std.testing.expectEqual(@as(usize, registry.len), families.len); +test "every registered family is listed for capabilities, with the wrappers it also reads" { + var n: usize = 0; + inline for (registry) |F| n += 1 + aliases(F).len; + try std.testing.expectEqual(n, families.len); try std.testing.expectEqualStrings("nemotron_h", families[0].model_type); + try std.testing.expectEqualStrings("nemotron_h_omni", families[1].model_type); + try std.testing.expect(reads(nemotron.native, "nemotron_h_omni") and !reads(nemotron.native, "qwen3_5")); } test "variables: unset or empty means the default, anything but a whole number is refused" { diff --git a/zig/src/native/cuda_alias.zig b/zig/src/native/cuda_alias.zig new file mode 100644 index 0000000000..0a9e9bc71a --- /dev/null +++ b/zig/src/native/cuda_alias.zig @@ -0,0 +1,13 @@ +//! The config model_types a CUDA family reads: its own, then its wrappers' (as Bolt's nemotron_h_omni). +const std = @import("std"); + +/// The other config model_types a family reads (`also_model_types`). +pub fn aliases(comptime F: type) []const []const u8 { + return if (@hasDecl(F, "also_model_types")) F.also_model_types else &.{}; +} + +pub fn reads(comptime F: type, model_type: []const u8) bool { + if (std.mem.eql(u8, model_type, F.model_type)) return true; + for (aliases(F)) |t| if (std.mem.eql(u8, model_type, t)) return true; + return false; +} diff --git a/zig/src/native/cuda_bound.zig b/zig/src/native/cuda_bound.zig index 40977dfa50..fe929ae849 100644 --- a/zig/src/native/cuda_bound.zig +++ b/zig/src/native/cuda_bound.zig @@ -18,6 +18,12 @@ pub fn vtable(comptime H: type, v: be.Backend.VTable) be.Backend.VTable { return x.inner.vtable.prefill_step.?(x.inner.ptr, s); } }.f else null, + .prefill_many = if (v.prefill_many != null) struct { + fn f(p: *anyopaque, ss: []const *lanes.Stream) anyerror!bool { + const x = H.bind(p); + return x.inner.vtable.prefill_many.?(x.inner.ptr, ss); + } + }.f else null, .first = struct { fn f(p: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { const x = H.bind(p); diff --git a/zig/src/native/cuda_memory.zig b/zig/src/native/cuda_memory.zig index 563b78bba1..6942f3e04d 100644 --- a/zig/src/native/cuda_memory.zig +++ b/zig/src/native/cuda_memory.zig @@ -102,6 +102,25 @@ pub fn cacheBytes(asked: ?f64, room: u64, streams: u64) u64 { return @min(room -| streams, want); } +/// Pinned host budget for kept prompt states: --prompt-cache-gib (0: none), else 8 GiB, at most half of MemAvailable. +pub fn hostCacheBytes(a: std.mem.Allocator, io: std.Io, asked: ?f64) u64 { + const want: u64 = if (asked) |g| (if (g > 0) std.math.lossyCast(u64, g * gib) else 0) else 8 << 30; + const text = procText(a, io, "/proc/meminfo") orelse return 0; + const m = meminfo(text) orelse return 0; + return hostShare(want, m.available); +} + +fn hostShare(want: u64, available: u64) u64 { + return @min(want, available / 2); +} + +test "host prompt cache: the ask, at most half of what is available" { + const g: u64 = 1 << 30; + try std.testing.expectEqual(8 * g, hostShare(8 * g, 64 * g)); + try std.testing.expectEqual(3 * g, hostShare(8 * g, 6 * g)); + try std.testing.expectEqual(@as(u64, 0), hostShare(0, 64 * g)); +} + /// Integrated GPUs require validated host counts; discrete devices use CUDA's counts. pub fn counts(unified: bool, card: MemInfo, text: ?[]const u8) error{HostMemoryUnavailable}!MemInfo { if (unified) return meminfo(text orelse return error.HostMemoryUnavailable) orelse error.HostMemoryUnavailable; diff --git a/zig/src/native/no_family.zig b/zig/src/native/no_family.zig new file mode 100644 index 0000000000..d81d7567ac --- /dev/null +++ b/zig/src/native/no_family.zig @@ -0,0 +1,2 @@ +//! Stands in for a CUDA family the build leaves out (-Dflashnext=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/native/no_family_glm.zig b/zig/src/native/no_family_glm.zig new file mode 100644 index 0000000000..2dcb07a292 --- /dev/null +++ b/zig/src/native/no_family_glm.zig @@ -0,0 +1,2 @@ +//! Stands in for the GLM CUDA family when the build leaves it out (-Dglm=false): nothing is registered. +pub const native = struct {}; diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 7c702ef432..63eebea817 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -35,6 +35,8 @@ pub const Input = struct { logprobs: ?u8 = null, /// ``/v1/completions``' ``prompt_logprobs``: rows for the prompt with this many alternatives; null: not asked. prompt_logprobs: ?u8 = null, + /// Image inputs (vision_inputs.extract), in prompt order; the rendered prompt's image pads expand to their tokens. + images: []const api.Image = &.{}, }; /// A streamed piece: content text (a string) or a delta object (reasoning or tool calls). @@ -90,9 +92,30 @@ pub const Prepared = struct { preparing: bool, }; +/// --tool-system: the server's instruction first in a tools request with no system message; others pass unchanged. +fn withToolSystem(srv: *Server, a: Allocator, input: Input) Allocator.Error!Input { + const text = srv.config.tool_system orelse return input; + if (input.tools.len == 0 or input.prompt != null or input.messages != .array) return input; + for (input.messages.array) |m| { + if (m != .object) continue; + const role = m.object.get("role") orelse continue; + if (role == .string and std.mem.eql(u8, role.string, "system")) return input; + } + const sys = try json.newObject(a); + try sys.put(a, "role", .{ .string = "system" }); + try sys.put(a, "content", .{ .string = text }); + const list = try a.alloc(Value, input.messages.array.len + 1); + list[0] = .{ .object = sys }; + @memcpy(list[1..], input.messages.array); + var out = input; + out.messages = .{ .array = list }; + return out; +} + /// Render ``input`` and run every check that can refuse it, before anything reaches the client or the engine. -pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepared { +pub fn prepare(srv: *Server, cx: *Cx, input_in: Input, gone: anytype) Failure!Prepared { const a = cx.a; + const input = try withToolSystem(srv, a, input_in); const io = srv.io; const received = nowNs(io); const f = input.fields; @@ -108,7 +131,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); - const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + var rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); + if (input.images.len > 0) rendered.history_len = 0; // image prompts keep no prefix for a later turn (as Python) if (gone.check()) return error.Cancelled; if (rendered.ids.len == 0) return cx.refuse("rendered prompt is empty"); const window: i64 = srv.info.context_window; @@ -119,7 +143,7 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (input.max_tokens != null and limit > room) return cx.fail(.context_length, "{s} {d} tokens, but the rendered prompt has {d} tokens and requests {d} reply tokens, which exceeds the context window. Reduce the prompt to at most {d} prompt tokens or request at most {d} reply tokens, including chat template and thinking tokens.", .{ errors.context_limit, window, n, limit, @max(0, window - limit), room }); limit = @min(limit, room); } - const system_len: usize = if (input.prompt != null) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); + const system_len: usize = if (input.prompt != null or input.images.len > 0) 0 else prompt_mod.systemPrefixLen(srv, cx, input.messages, input.tools, rendered.ids, thinking, effort); var shared: std.ArrayList(u32) = .empty; if (system_len > 0) for ([_]i64{ @as(i64, @intCast(system_len)) - 2048, @as(i64, @intCast(system_len)) - 512, @intCast(system_len) }) |cut| { if (cut >= 512) try shared.append(a, @intCast(cut)); @@ -137,7 +161,9 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .history_len = @intCast(rendered.history_len), .rewind_len = @intCast(rendered.rewind_len), .shared_prefixes = shared.items, - .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), // a cut before the conversation's own text: fresh sessions resume their whole harness + .images = input.images, + // a cut just before the conversation's own text: fresh sessions resume their whole harness + .chunks = try chunk_plan.withCut(a, try srv.chunks.starts(a, rendered.ids), if (srv.chunks.step > 0) @intCast(@max(system_len, 1) - 1) else 0, rendered.ids.len, srv.chunks.min_chunk), .decode_spans = try prompt_mod.replySpans(srv, cx, rendered.ids), // from the tokens alone, for every request: cached states are keyed by tokens .tools_json = if (input.tools.len > 0) try json.stringify(a, .{ .array = @constCast(input.tools) }, .{ .ascii = false }) else "", .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, @@ -291,7 +317,7 @@ const Generation = struct { return !g.ignore_eos and std.mem.indexOfScalar(u32, g.srv.eos, t) != null; } - /// What closes a call the model's end token left open (``tool_parse.closeCall``); nothing when a stop string, the length or ``ignore_eos`` ended the reply instead. + /// What closes a call the end token left open (``tool_parse.closeCall``); nothing for stop, length or ignore_eos. fn closeCall(g: *const Generation, text: []const u8) Allocator.Error![]const u8 { const t = g.collected.items; if (g.tools.len == 0 or t.len == 0 or !g.eos(t[t.len - 1])) return ""; @@ -371,7 +397,7 @@ const Generation = struct { }; } - /// Ends the request and waits for the engine; one it ended unfinished counts as a disconnect, as the Mac scheduler counts it. + /// Ends the request and waits for the engine; one ended unfinished counts as a disconnect, as on the Mac. fn cancel(g: *Generation) void { if (!g.engine_done) g.srv.engine.cancel(g.id); g.drain(); diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..852b3730b7 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -49,15 +49,16 @@ pub const flags = [_]Flag{ .{ .name = "--api-key-file", .native = true }, .{ .name = "--metrics-open", .kind = .store_true, .native = true }, .{ .name = "--dashboard", .kind = .store_true, .native = true }, - .{ .name = "--vision", .kind = .store_true }, - .{ .name = "--vision-urls", .kind = .store_true }, + .{ .name = "--vision", .kind = .store_true, .native = true }, + .{ .name = "--vision-urls", .kind = .store_true, .native = true }, .{ .name = "--vision-offload", .kind = .store_true }, - .{ .name = "--vision-max-images" }, - .{ .name = "--vision-image-tokens" }, + .{ .name = "--vision-max-images", .native = true }, + .{ .name = "--vision-image-tokens", .native = true }, .{ .name = "--context", .native = true }, .{ .name = "--speed-up", .native = !cuda_build }, .{ .name = "--max-tokens", .native = true }, .{ .name = "--temperature", .native = true }, + .{ .name = "--tool-system", .native = true }, .{ .name = "--top-p", .native = true }, .{ .name = "--top-k", .native = true }, .{ .name = "--min-p", .native = true }, @@ -99,7 +100,7 @@ pub const flags = [_]Flag{ .{ .name = "--master-port", .native = served("--tp") }, // the GPU engine's policy switches as key=value pairs, repeatable .{ .name = "--policy", .kind = .append, .native = served("--policy") }, - .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" } }, + .{ .name = "--kv-dtype", .choices = &.{ "bf16", "int8", "int4" }, .native = served("--kv-dtype"), .native_values = &.{ "int8", "int4" } }, .{ .name = "--prefill-fp8", .kind = .store_true }, .{ .name = "--no-prefill-fp8", .kind = .store_true }, .{ .name = "--precision", .choices = &.{ "checkpoint", "full" } }, @@ -140,6 +141,11 @@ pub const Args = struct { learn_gib: f64 = 32, // disk for learned states, every model and build together max_tokens: i64 = 4096, temperature: ?f64 = null, + vision: bool = false, // accept image inputs (an engine whose family has a vision tower) + vision_urls: bool = false, // also public HTTPS image and video URLs + vision_max_images: i64 = 4, + vision_image_tokens: i64 = 4096, // visual tokens a request's images share + tool_system: ?[]const u8 = null, // an instruction added as the system message of tool requests that send none top_p: ?f64 = null, top_k: ?i64 = null, min_p: ?f64 = null, @@ -170,6 +176,7 @@ pub const Args = struct { master_port: u16 = 29551, policy: []const u8 = "", // every --policy joined by commas cortheon: bool = false, // Cortheon's program, research before building, and cortheon_search/cortheon_fetch run by the server + kv_dtype: []const u8 = "int8", // --kv-dtype (CUDA): the attention cache's codes, int8 or int4 }; /// A usage error's message (argparse's ``error:`` line); the caller exits 2. @@ -265,7 +272,7 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa } else if (is(name, "--learn-gib")) { out.learn = true; out.learn_gib = try gib(u, a, name, v); - } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { + } else if (is(name, "--max-tokens")) out.max_tokens = try int(u, a, name, v) else if (is(name, "--temperature")) out.temperature = try float(u, a, name, v) else if (is(name, "--tool-system")) out.tool_system = v else if (is(name, "--vision")) out.vision = true else if (is(name, "--vision-urls")) out.vision_urls = true else if (is(name, "--vision-max-images")) out.vision_max_images = try int(u, a, name, v) else if (is(name, "--vision-image-tokens")) out.vision_image_tokens = try int(u, a, name, v) else if (is(name, "--top-p")) out.top_p = try float(u, a, name, v) else if (is(name, "--top-k")) out.top_k = try int(u, a, name, v) else if (is(name, "--min-p")) out.min_p = try float(u, a, name, v) else if (is(name, "--thinking")) out.thinking = true else if (is(name, "--no-thinking")) out.thinking = false else if (is(name, "--reasoning-effort")) out.reasoning_effort = v else if (is(name, "--thinking-budget")) out.thinking_budget = try int(u, a, name, v) else if (is(name, "--loop-guard")) out.loop_guard = true else if (is(name, "--drafter")) out.drafter = v else if (is(name, "--drafter-bits")) out.drafter_bits = @intCast(try int(u, a, name, v)) else if (is(name, "--no-drafts")) out.no_drafts = true else if (is(name, "--keep-warm")) out.keep_warm = try int(u, a, name, v) else if (is(name, "--compact-at")) { if (std.mem.eql(u8, v, "auto")) out.compact_auto = true else { const f = try float(u, a, name, v); if (!(f > 0 and f <= 1)) return fail(u, a, "argument --compact-at: expected auto or a fraction in (0, 1]: '{s}'", .{v}); @@ -282,9 +289,19 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa fn cudaFlag(a: Allocator, out: *Args, name: []const u8, v: []const u8, u: *Usage) error{ Usage, OutOfMemory }!bool { if (std.mem.eql(u8, name, "--device")) { out.device = std.math.cast(u32, try int(u, a, name, v)) orelse return fail(u, a, "argument --device: a GPU ordinal from 0: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--kv-dtype")) { + out.kv_dtype = v; } else if (std.mem.eql(u8, name, "--segments")) { const n = try int(u, a, name, v); out.segments = if (n >= 1) @intCast(@min(n, std.math.maxInt(u32))) else return fail(u, a, "argument --segments: a count from 1: '{s}'", .{v}); + } else if (std.mem.eql(u8, name, "--tp")) { + out.tp = if (std.mem.eql(u8, v, "2")) 2 else 1; + } else if (std.mem.eql(u8, name, "--rank")) { + out.rank = if (std.mem.eql(u8, v, "1")) 1 else 0; + } else if (std.mem.eql(u8, name, "--master")) { + out.master = v; + } else if (std.mem.eql(u8, name, "--master-port")) { + out.master_port = std.math.cast(u16, try int(u, a, name, v)) orelse return fail(u, a, "argument --master-port: a TCP port: '{s}'", .{v}); } else return false; return true; } @@ -362,6 +379,8 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .mtp_drafts = args.mtp_drafts, .mtp_confidence = args.mtp_confidence, .keep = args.checkpoint_slots, + .vision = args.vision, + .kv_bits = if (std.mem.eql(u8, args.kv_dtype, "int4")) 4 else 8, .tp = args.tp, .rank = args.rank, .master = args.master, diff --git a/zig/src/server/main.zig b/zig/src/server/main.zig index 4b85fca8cd..18a6d3a5f7 100644 --- a/zig/src/server/main.zig +++ b/zig/src/server/main.zig @@ -78,6 +78,10 @@ pub fn main(init: std.process.Init) !u8 { } orelse return fail(problem); defer up.text.deinit(); defer up.engine.close(up.engine.ctx); + if (up.engine.follow) |follow| { // rank 1 of two: rank 0 serves, this rank replays its engine calls + follow(up.engine.ctx) catch |e| return fail(@errorName(e)); + return 0; + } return serve.run(gpa, io, args, .{ .engine = up.engine.engine, .text = up.text.text(), diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig new file mode 100644 index 0000000000..c9ef04847f --- /dev/null +++ b/zig/src/server/media_fetch.zig @@ -0,0 +1,330 @@ +//! Public HTTPS media for --vision-urls as vision/images_http.py fetches it: checked hosts, TLS, 3 redirects, limits. + +const std = @import("std"); +const Io = std.Io; +const net = std.Io.net; + +pub const Error = error{ OutOfMemory, Media }; + +/// The failure's words for the client (the request is refused with them). +pub const Failure = struct { text: []const u8 = "" }; + +pub const max_url_chars = 4096; +pub const max_redirects = 3; +pub const timeout_s = 10; // one download +pub const total_s = 30; // a request's downloads together + +fn fail(f: *Failure, text: []const u8) Error { + f.text = text; + return error.Media; +} + +/// Python ipaddress is_global, less multicast, reserved, 192.0.0.0/24, 168.63.129.16 and IPv6 embedding an IPv4 one. +pub fn publicIp(a: net.IpAddress) bool { + switch (a) { + .ip4 => |v| { + const b = v.bytes; + const in = struct { + fn net4(x: [4]u8, base: [4]u8, bits: u5) bool { + const xi = std.mem.readInt(u32, &x, .big); + const bi = std.mem.readInt(u32, &base, .big); + const mask: u32 = if (bits == 0) 0 else ~@as(u32, 0) << @intCast(32 - @as(u6, bits)); + return xi & mask == bi & mask; + } + }.net4; + const blocked = [_]struct { [4]u8, u5 }{ + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, + }; + for (blocked) |blk| if (in(b, blk[0], blk[1])) return false; + if (std.mem.eql(u8, &b, &.{ 168, 63, 129, 16 })) return false; + return true; + }, + .ip6 => |v| { + const b = v.bytes; + // 2000::/3 only, less 2001:db8::/32, the 2001::/23 IETF block (Teredo among it) and 6to4 2002::/16 + if (b[0] & 0xe0 != 0x20) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; + if (b[0] == 0x20 and b[1] == 0x02) return false; + return true; + }, + } +} + +const Url = struct { host: []const u8, target: []const u8 }; + +/// grammar of images_http._url: https, a host, port 443, no credentials, fragment, backslash, whitespace or controls. +fn parseUrl(a: std.mem.Allocator, value: []const u8, f: *Failure) Error!Url { + if (value.len > max_url_chars) return fail(f, "image URL is too long or contains whitespace/control characters"); + for (value) |c| if (c <= 32 or c == 127) return fail(f, "image URL is too long or contains whitespace/control characters"); + const bad = "image URL must be HTTPS on port 443, without credentials or a fragment"; + if (!std.ascii.startsWithIgnoreCase(value, "https://")) return fail(f, bad); + if (std.mem.indexOfScalar(u8, value, '\\') != null or std.mem.indexOfScalar(u8, value, '#') != null) return fail(f, bad); + const rest = value["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const authority = rest[0..end]; + if (std.mem.indexOfScalar(u8, authority, '@') != null) return fail(f, bad); + var host = authority; + if (std.mem.lastIndexOfScalar(u8, authority, ':')) |c| { + if (authority.len > 0 and authority[0] == '[') return fail(f, bad); // a literal IPv6 host: not a name + if (!std.mem.eql(u8, authority[c + 1 ..], "443") and authority[c + 1 ..].len > 0) return fail(f, bad); + host = authority[0..c]; + } + if (host.len == 0 or std.mem.indexOfScalar(u8, host, '%') != null) return fail(f, bad); + for (host) |c| if (c >= 128) return fail(f, bad); // IDNA names: their ASCII (xn--) form + const lower = try std.ascii.allocLowerString(a, std.mem.trimEnd(u8, host, ".")); + for ([_][]const u8{ "localhost", "metadata.google.internal", "instance-data" }) |n| if (std.mem.eql(u8, lower, n)) return fail(f, "image URLs must use public internet hosts"); + // the path and query, percent-quoted as urllib.parse.quote(safe=...) leaves them + var t: std.ArrayList(u8) = .empty; + var tail = rest[end..]; + if (tail.len == 0 or tail[0] == '?') try t.append(a, '/'); + const safe_path = "/%:@!$&'()*+,;=-._~"; + var in_query = false; + for (tail) |c| { + if (c == '?' and !in_query) { + in_query = true; + try t.append(a, '?'); + continue; + } + const safe = std.ascii.isAlphanumeric(c) or std.mem.indexOfScalar(u8, safe_path, c) != null or (in_query and c == '?'); + if (safe) try t.append(a, c) else try t.print(a, "%{X:0>2}", .{c}); + } + tail = t.items; + return .{ .host = lower, .target = tail }; +} + +/// The location a redirect names, resolved against ``base`` (urljoin for absolute and absolute-path forms). +fn join(a: std.mem.Allocator, base: []const u8, location: []const u8) ![]const u8 { + if (std.ascii.startsWithIgnoreCase(location, "https://") or std.ascii.startsWithIgnoreCase(location, "http://")) return a.dupe(u8, location); + if (std.mem.startsWith(u8, location, "//")) return std.fmt.allocPrint(a, "https:{s}", .{location}); + const rest = base["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const origin = base[0 .. "https://".len + end]; + if (std.mem.startsWith(u8, location, "/")) return std.fmt.allocPrint(a, "{s}{s}", .{ origin, location }); + const path_end = std.mem.indexOfScalar(u8, rest[end..], '?') orelse rest.len - end; + const path = rest[end..][0..path_end]; + const dir = if (std.mem.lastIndexOfScalar(u8, path, '/')) |s| path[0 .. s + 1] else "/"; + return std.fmt.allocPrint(a, "{s}{s}{s}", .{ origin, dir, location }); +} + +fn remaining(io: Io, deadline: Io.Timestamp, f: *Failure) Error!i64 { + const left = Io.Clock.awake.now(io).durationTo(deadline).toMilliseconds(); + if (left <= 0) return fail(f, "image download timed out"); + return left; +} + +fn setTimeouts(handle: net.Socket.Handle, ms: i64) void { + const tv: std.posix.timeval = .{ .sec = @intCast(@divTrunc(ms, 1000)), .usec = @intCast(@mod(ms, 1000) * 1000) }; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.RCVTIMEO, std.mem.asBytes(&tv)) catch {}; + std.posix.setsockopt(handle, std.posix.SOL.SOCKET, std.posix.SO.SNDTIMEO, std.mem.asBytes(&tv)) catch {}; +} + +/// A TCP connection to the checked address, connect, sends and reads bounded by ``ms`` (SO_SNDTIMEO bounds connect). +fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + const posix = std.posix; + const linux = std.os.linux; + const family: u32 = switch (addr) { + .ip4 => posix.AF.INET, + .ip6 => posix.AF.INET6, + }; + const rc = linux.socket(family, posix.SOCK.STREAM | posix.SOCK.CLOEXEC, 0); + if (linux.errno(rc) != .SUCCESS) return fail(f, "image download failed or timed out"); + const fd: posix.fd_t = @intCast(rc); + errdefer _ = linux.close(fd); + setTimeouts(fd, ms); + var sa4: posix.sockaddr.in = undefined; + var sa6: posix.sockaddr.in6 = undefined; + const r = switch (addr) { + .ip4 => |v| blk: { + sa4 = .{ .port = std.mem.nativeToBig(u16, 443), .addr = @bitCast(v.bytes) }; + break :blk linux.connect(fd, @ptrCast(&sa4), @sizeOf(posix.sockaddr.in)); + }, + .ip6 => |v| blk: { + sa6 = .{ .port = std.mem.nativeToBig(u16, 443), .flowinfo = 0, .addr = v.bytes, .scope_id = 0 }; + break :blk linux.connect(fd, @ptrCast(&sa6), @sizeOf(posix.sockaddr.in6)); + }, + }; + if (linux.errno(r) != .SUCCESS) return fail(f, "image download failed or timed out"); + return .{ .socket = .{ .handle = fd, .address = addr } }; +} + +/// The bytes of ``url`` (and its media type, one of ``media_types``), at most ``max_bytes``, before ``deadline``. +pub fn fetch(a: std.mem.Allocator, io: Io, url_in: []const u8, max_bytes: usize, deadline_in: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!struct { data: []u8, media: []const u8 } { + var url = url_in; + const deadline = blk: { + const one = Io.Clock.awake.now(io).addDuration(.fromSeconds(timeout_s)); + break :blk if (one.nanoseconds < deadline_in.nanoseconds) one else deadline_in; + }; + var hop: usize = 0; + while (true) : (hop += 1) { + const u = try parseUrl(a, url, f); + const got = try request(a, io, u, max_bytes, deadline, media_types, f); + switch (got) { + .body => |b| return .{ .data = b.data, .media = b.media }, + .redirect => |loc| { + if (hop == max_redirects) return fail(f, "image download has too many redirects"); + url = join(a, url, loc) catch return error.OutOfMemory; + }, + } + } +} + +const Got = union(enum) { body: struct { data: []u8, media: []const u8 }, redirect: []const u8 }; + +var bundle_lock: Io.RwLock = .init; +var bundle: std.crypto.Certificate.Bundle = .empty; +var bundle_loaded = false; + +fn request(a: std.mem.Allocator, io: Io, u: Url, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + // every address the name resolves to must be public; the connection goes to the first of them + const name = net.HostName.init(u.host) catch return fail(f, "image host could not be resolved"); + var results: [32]net.HostName.LookupResult = undefined; + var q: Io.Queue(net.HostName.LookupResult) = .init(&results); + name.lookup(io, &q, .{ .port = 443 }) catch return fail(f, "image host could not be resolved"); + var first: ?net.IpAddress = null; + var count: usize = 0; + while (q.getOneUncancelable(io)) |r| switch (r) { + .address => |addr| { + if (!publicIp(addr)) return fail(f, "image URLs must resolve only to public internet addresses"); + if (first == null) first = addr; + count += 1; + }, + .canonical_name => {}, + } else |_| {} + const addr = first orelse return fail(f, "image host could not be resolved"); + _ = try remaining(io, deadline, f); + // the CA bundle, read once + { + bundle_lock.lockUncancelable(io); + defer bundle_lock.unlock(io); + if (!bundle_loaded) { + bundle.rescan(std.heap.page_allocator, io, Io.Clock.real.now(io)) catch return fail(f, "image download failed: no CA certificates on the server"); + bundle_loaded = true; + } + } + const stream = try connect(addr, try remaining(io, deadline, f), f); + defer stream.close(io); + const tls_len = std.crypto.tls.Client.min_buffer_len; + const bufs = try a.alloc(u8, 4 * tls_len); + var sr = stream.reader(io, bufs[0..tls_len]); + var sw = stream.writer(io, bufs[tls_len .. 2 * tls_len]); + var entropy: [std.crypto.tls.Client.Options.entropy_len]u8 = undefined; + io.random(&entropy); + var tls = std.crypto.tls.Client.init(&sr.interface, &sw.interface, .{ + .host = .{ .explicit = u.host }, + .ca = .{ .bundle = .{ .gpa = std.heap.page_allocator, .io = io, .lock = &bundle_lock, .bundle = &bundle } }, + .read_buffer = bufs[2 * tls_len .. 3 * tls_len], + .write_buffer = bufs[3 * tls_len ..], + .entropy = &entropy, + .realtime_now = Io.Clock.real.now(io), + .allow_truncation_attacks = true, // HTTP's own lengths and framing end the body + }) catch return fail(f, "image download failed: TLS could not be verified for the host"); + // the request + const accept = std.mem.join(a, ", ", media_types) catch return error.OutOfMemory; + tls.writer.print("GET {s} HTTP/1.1\r\nHost: {s}\r\nAccept: {s}\r\nAccept-Encoding: identity\r\nUser-Agent: TensorFold-native\r\nConnection: close\r\n\r\n", .{ u.target, u.host, accept }) catch return fail(f, "image download failed or timed out"); + tls.writer.flush() catch return fail(f, "image download failed or timed out"); + sw.interface.flush() catch return fail(f, "image download failed or timed out"); + const r = &tls.reader; + // status and headers + const status_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + if (status_line.len < 12 or !std.mem.startsWith(u8, status_line, "HTTP/1.")) return fail(f, "image download failed"); + const status = std.fmt.parseInt(u16, status_line[9..12], 10) catch return fail(f, "image download failed"); + var location: ?[]const u8 = null; + var media: []const u8 = ""; + var length: ?usize = null; + var chunked = false; + var encoded = false; + var header_bytes: usize = 0; + while (true) { + const line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + header_bytes += line.len; + if (header_bytes > 64 * 1024) return fail(f, "image download failed"); + const l = std.mem.trimEnd(u8, line, "\r\n"); + if (l.len == 0) break; + const colon = std.mem.indexOfScalar(u8, l, ':') orelse continue; + const key = std.mem.trim(u8, l[0..colon], " \t"); + const val = std.mem.trim(u8, l[colon + 1 ..], " \t"); + if (std.ascii.eqlIgnoreCase(key, "location")) location = try a.dupe(u8, val); + if (std.ascii.eqlIgnoreCase(key, "content-type")) { + const semi = std.mem.indexOfScalar(u8, val, ';') orelse val.len; + media = try std.ascii.allocLowerString(a, std.mem.trim(u8, val[0..semi], " \t")); + } + if (std.ascii.eqlIgnoreCase(key, "content-length")) { + for (val) |c| if (!std.ascii.isDigit(c)) return fail(f, "image response exceeds the encoded byte limit"); + length = std.fmt.parseInt(usize, val, 10) catch return fail(f, "image response exceeds the encoded byte limit"); + } + if (std.ascii.eqlIgnoreCase(key, "transfer-encoding")) chunked = std.mem.indexOf(u8, try std.ascii.allocLowerString(a, val), "chunked") != null; + if (std.ascii.eqlIgnoreCase(key, "content-encoding") and !std.ascii.eqlIgnoreCase(val, "identity")) encoded = true; + } + switch (status) { + 301, 302, 303, 307, 308 => return .{ .redirect = location orelse return fail(f, "image redirect has no destination") }, + 200 => {}, + else => return fail(f, try std.fmt.allocPrint(a, "image download returned HTTP {d}", .{status})), + } + if (encoded) return fail(f, "compressed HTTP image responses are unsupported"); + for (media_types) |m| { + if (std.mem.eql(u8, m, media)) break; + } else return fail(f, if (std.mem.startsWith(u8, media_types[0], "image/")) "image URL content type must be JPEG, PNG or WebP" else try std.fmt.allocPrint(a, "media URL content type must be one of {s}", .{accept})); + if (length) |n| if (n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + var data: std.ArrayList(u8) = .empty; + if (chunked) { + while (true) { + _ = try remaining(io, deadline, f); + const size_line = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed or timed out"); + const semi = std.mem.indexOfScalar(u8, size_line, ';') orelse size_line.len; + const hex = std.mem.trim(u8, size_line[0..semi], " \t\r\n"); + const n = std.fmt.parseInt(usize, hex, 16) catch return fail(f, "image download failed"); + if (n == 0) break; + if (data.items.len + n > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + const piece = r.take(n) catch return fail(f, "image download failed or timed out"); + try data.appendSlice(a, piece); + _ = r.takeDelimiterInclusive('\n') catch return fail(f, "image download failed"); + } + } else { + while (length == null or data.items.len < length.?) { + _ = try remaining(io, deadline, f); + const want = if (length) |n| @min(n - data.items.len, 64 * 1024) else 64 * 1024; + const piece = r.peekGreedy(1) catch |e| switch (e) { + error.EndOfStream => break, + else => return fail(f, "image download failed or timed out"), + }; + const take = @min(piece.len, want); + if (data.items.len + take > max_bytes) return fail(f, "image response exceeds the encoded byte limit"); + try data.appendSlice(a, piece[0..take]); + r.toss(take); + } + if (length) |n| if (data.items.len != n) return fail(f, "image download failed or timed out"); + } + _ = try remaining(io, deadline, f); + return .{ .body = .{ .data = data.items, .media = media } }; +} + +test "public addresses as Python's ipaddress checks them" { + const ip4 = struct { + fn of(b: [4]u8) net.IpAddress { + return .{ .ip4 = .{ .bytes = b, .port = 443 } }; + } + }.of; + try std.testing.expect(publicIp(ip4(.{ 8, 8, 8, 8 }))); + try std.testing.expect(publicIp(ip4(.{ 151, 101, 1, 69 }))); + for ([_][4]u8{ .{ 10, 1, 2, 3 }, .{ 127, 0, 0, 1 }, .{ 169, 254, 169, 254 }, .{ 172, 20, 0, 1 }, .{ 192, 168, 1, 1 }, .{ 100, 100, 0, 1 }, .{ 0, 0, 0, 0 }, .{ 224, 0, 0, 1 }, .{ 255, 255, 255, 255 }, .{ 168, 63, 129, 16 }, .{ 192, 0, 0, 9 } }) |b| + try std.testing.expect(!publicIp(ip4(b))); +} + +test "URLs as images_http._url takes them" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Failure = .{}; + try std.testing.expectError(error.Media, parseUrl(a, "https://example.com/a b/c.png", &f)); + const ok = try parseUrl(a, "https://example.com:443/img/c.png?x=%20&y=ä", &f); + try std.testing.expectEqualStrings("example.com", ok.host); + try std.testing.expectEqualStrings("/img/c.png?x=%20&y=%C3%A4", ok.target); + for ([_][]const u8{ "http://example.com/x.png", "https://user@example.com/x", "https://example.com:8443/x", "https://localhost/x", "https://example.com/x#frag", "https://[::1]/x" }) |bad| + try std.testing.expectError(error.Media, parseUrl(a, bad, &f)); + try std.testing.expectEqualStrings("https://e.com/b/c.png", try join(a, "https://e.com/b/a.png", "c.png")); + try std.testing.expectEqualStrings("https://e.com/z", try join(a, "https://e.com/b/a.png", "/z")); +} diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..aaf23610ae 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -127,7 +127,12 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (srv.config.request_log) |path| request_log.append(cx.a, path, body); var input: chat.Input = .{ .fields = undefined }; if (is_chat) { - input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + var msgs = body.get("messages"); + if (try @import("vision_inputs.zig").extract(srv, cx, msgs)) |x| { // images: decoded, prepared, their parts as text + msgs = x.messages; + input.images = x.images; + } + input.messages = try messages.normalize(cx, msgs, "system", srv.needs_user_after_tool); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..f48b0a76dd 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -39,6 +39,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t error.OutOfMemory => return error.OutOfMemory, error.Template => return cx.fail(.server, "{s}", .{problem}), }; + if (!thinking) ids = try withoutEffortLine(srv, cx, ids); if (!thinking and generation and ids.len > 0) { const last = try srv.text.decode(cx.a, ids[ids.len - 1 ..]); if (std.mem.eql(u8, @import("reply_text.zig").pyStrip(last), "")) { @@ -52,6 +53,17 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } +/// GLM-5.3's effort system line, dropped with thinking off as glm5_next/prompts.py does; its ids are its own segment. +fn withoutEffortLine(srv: *Server, cx: *Cx, ids: []u32) errors.Refused![]u32 { + const line = srv.text.encode(cx.a, "<|system|>Reasoning Effort: Max", false) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Template => return ids, + }; + if (line.len < 2) return ids; + const at = std.mem.indexOf(u32, ids, line) orelse return ids; + return try std.mem.concat(cx.a, u32, &.{ ids[0..at], ids[at + line.len ..] }); +} + /// The template's generation prompt (thinking on or off): what a probe conversation renders past its history. fn generationPrompt(srv: *Server, cx: *Cx, thinking: bool) errors.Refused![]const u32 { const user = try json.newObject(cx.a); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..1bebd524e3 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -109,6 +109,10 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .model_dir = s.model_dir, .model_ids = ids.items, .default_max_tokens = args.max_tokens, + .vision_max_images = @intCast(@max(args.vision_max_images, 1)), + .vision_urls = args.vision_urls, + .vision_image_tokens = @intCast(std.math.clamp(args.vision_image_tokens, 1, 65536)), + .tool_system = if (args.tool_system) |t| (if (std.mem.trim(u8, t, " \t\r\n").len > 0) std.mem.trim(u8, t, " \t\r\n") else null) else null, .enable_thinking = args.thinking, .reasoning_effort = args.reasoning_effort, .thinking_budget = args.thinking_budget, diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..3b5f0a1746 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -36,6 +36,10 @@ pub const Config = struct { /// The served name first, then aliases, without repeats. model_ids: []const []const u8, default_max_tokens: i64 = 4096, + tool_system: ?[]const u8 = null, // --tool-system (trimmed; null when empty) + vision_max_images: usize = 4, // --vision-max-images + vision_urls: bool = false, // --vision-urls: public HTTPS media URLs are fetched (media_fetch.zig) + vision_image_tokens: u64 = 4096, // --vision-image-tokens: the visual tokens a request's images share enable_thinking: bool = true, reasoning_effort: ?[]const u8 = null, thinking_budget: i64 = 0, @@ -135,6 +139,7 @@ pub const Server = struct { if (srv.keepalive) |k| k.stop(); // before the engine's queue goes away if (srv.grammars) |g| g.close(); srv.metrics.deinit(); + if (srv.grammars) |g| g.close(); srv.arena.deinit(); srv.store.deinit(); srv.slide.deinit(); @@ -211,7 +216,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 grammar compiler: tokenizer.json, config.json's vocab_size (text_config or llm_config first), the eos ids. fn grammarCompiler(srv: *Server) ?*api.grammar.Compiler { srv.grammar_mutex.lockUncancelable(srv.io); defer srv.grammar_mutex.unlock(srv.io); @@ -227,7 +232,7 @@ pub const Server = struct { .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| { + for ([_]?json.Value{ cfg.get("text_config"), cfg.get("llm_config"), cfg }) |part| if (vocab == null) if (part) |p| if (p.get("vocab_size")) |v| if (v.int64()) |n| { vocab = @intCast(n); }; if (vocab == null) break :blk "structured output needs the checkpoint's config.json vocab_size"; diff --git a/zig/src/server/vision_inputs.zig b/zig/src/server/vision_inputs.zig new file mode 100644 index 0000000000..2477926109 --- /dev/null +++ b/zig/src/server/vision_inputs.zig @@ -0,0 +1,350 @@ +//! A chat request's image_url and video_url parts, prepared and replaced by the text Python's frontend expands them to. + +const std = @import("std"); +const media_fetch = @import("media_fetch.zig"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const Server = @import("server.zig").Server; +const Value = json.Value; +const Cx = errors.Cx; +const qi = api.qwen_image; + +const vision_start = "<|vision_start|>"; +const vision_end = "<|vision_end|>"; +const max_video_bytes: usize = 16 * 1024 * 1024; +const max_videos = 2; +const max_bytes_one: usize = 10 * 1024 * 1024; +const max_bytes_all: usize = 20 * 1024 * 1024; +const max_pixels_all: u64 = 32 * 1024 * 1024; +const max_tokens_one: u64 = 4096; + +pub const Extracted = struct { messages: Value, images: []api.Image }; + +/// The image URL of a content part (``image_url``, ``image`` or ``input_image``) and its ``detail``; else null. +const Kind = enum { image, video }; + +fn mediaPart(part: Value) ?struct { kind: Kind, url: ?[]const u8, detail: []const u8 } { + if (part != .object) return null; + const ty = part.get("type"); + const is_video = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "video_url") or std.mem.eql(u8, ty.?.string, "video"))) or + part.get("video_url") != null; + if (is_video) { + var vurl: ?[]const u8 = null; + if (part.get("video_url")) |vu| switch (vu) { + .string => |x| vurl = x, + .object => if (vu.get("url")) |u| if (u == .string) { + vurl = u.string; + }, + else => {}, + }; + if (vurl == null) if (part.get("video")) |v| if (v == .string) { + vurl = v.string; + }; + return .{ .kind = .video, .url = vurl, .detail = "auto" }; + } + const is_image = (ty != null and ty.? == .string and (std.mem.eql(u8, ty.?.string, "image_url") or std.mem.eql(u8, ty.?.string, "image") or + std.mem.eql(u8, ty.?.string, "input_image"))) or part.get("image_url") != null or part.get("image") != null; + if (!is_image) return null; + var url: ?[]const u8 = null; + var detail: []const u8 = "auto"; + if (part.get("image_url")) |iu| switch (iu) { + .string => |s| url = s, + .object => { + if (iu.get("url")) |u| if (u == .string) { + url = u.string; + }; + if (iu.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + }, + else => {}, + }; + if (url == null) if (part.get("image")) |im| if (im == .string) { + url = im.string; + }; + if (part.get("detail")) |d| if (d == .string) { + detail = d.string; + }; + return .{ .kind = .image, .url = url, .detail = detail }; +} + +fn hasImages(messages: Value) bool { + if (messages != .array) return false; + for (messages.array) |m| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + for (c.array) |p| if (mediaPart(p) != null) return true; + } + return false; +} + +const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; +const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; + +/// The bytes of a ``data:`` URL, or (--vision-urls) a public HTTPS URL fetched before ``deadline`` (images.py). +fn dataBytes(srv: *Server, cx: *Cx, url: []const u8, limit: usize, kind: Kind, deadline: std.Io.Timestamp) errors.Refused![]u8 { + if (!std.mem.startsWith(u8, url, "data:")) { + if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse(if (kind == .image) "images require data URLs or public HTTPS URLs" else "videos require data URLs or public HTTPS URLs"); + if (!srv.config.vision_urls) return cx.refuse("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls"); + if (url.len > media_fetch.max_url_chars) return cx.refuse("image URL is too long"); + var f: media_fetch.Failure = .{}; + const got = media_fetch.fetch(cx.a, srv.io, url, limit, deadline, if (kind == .image) &image_media else &video_media, &f) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Media => return cx.fail(.request, "{s}", .{f.text}), + }; + return got.data; + } + const comma = std.mem.indexOfScalar(u8, url, ',') orelse return cx.refuse("an image data: URL needs a comma before its bytes"); + if (!std.mem.endsWith(u8, url[0..comma], ";base64")) return cx.refuse("an image data: URL must be base64-encoded"); + const b64 = url[comma + 1 ..]; + const dec = std.base64.standard.Decoder; + const n = dec.calcSizeForSlice(b64) catch return cx.refuse("an image's base64 bytes are invalid"); + if (n > limit) return cx.refuse(if (kind == .image) "an image exceeds the 10 MiB limit (20 MiB for a request's media together)" else "a video exceeds the 16 MiB limit (20 MiB for a request's media together)"); + const out = try cx.a.alloc(u8, n); + dec.decode(out, b64) catch return cx.refuse("an image's base64 bytes are invalid"); + return out; +} + +/// A request's images prepared, its image parts as the template's image text; null without images. Needs --vision. +pub fn extract(srv: *Server, cx: *Cx, messages: ?Value) errors.Refused!?Extracted { + const list = messages orelse return null; + if (!hasImages(list)) return null; + if (!srv.info.vision) return cx.refuse("this server accepts text parts only; image, audio and video inputs are unsupported"); + // every media part, in prompt order, with where its text goes + const Slot = struct { kind: Kind, url: ?[]const u8, detail: []const u8, msg: usize, part: usize }; + var slots: std.ArrayList(Slot) = .empty; + for (list.array, 0..) |m, i| { + if (m != .object) continue; + const c = m.get("content") orelse continue; + if (c != .array) continue; + const role = m.get("role"); + for (c.array, 0..) |p, j| { + const im = mediaPart(p) orelse continue; + if (role != null and role.? == .string and (std.mem.eql(u8, role.?.string, "system") or std.mem.eql(u8, role.?.string, "developer"))) + return cx.refuse("a system message cannot contain images or videos"); + try slots.append(cx.a, .{ .kind = im.kind, .url = im.url, .detail = im.detail, .msg = i, .part = j }); + } + } + var n_images: usize = 0; + var n_videos: usize = 0; + for (slots.items) |sl| switch (sl.kind) { + .image => n_images += 1, + .video => n_videos += 1, + }; + if (n_images > srv.config.vision_max_images) return cx.fail(.request, "a request may carry at most {d} images; this one has {d}", .{ srv.config.vision_max_images, n_images }); + if (n_videos > max_videos) return cx.fail(.request, "a request may carry at most {d} videos; this one has {d}", .{ max_videos, n_videos }); + const budget: u64 = srv.config.vision_image_tokens; + if (n_images > 0 and budget < n_images) return cx.refuse("the image count exceeds the visual-token budget"); + const per: u64 = if (n_images > 0) @min(budget / n_images, max_tokens_one) else 0; + var images: std.ArrayList(api.Image) = .empty; + const texts = try cx.a.alloc([]const u8, slots.items.len); + var total_bytes: usize = 0; + var total_pixels: u64 = 0; + const deadline = std.Io.Clock.awake.now(srv.io).addDuration(.fromSeconds(media_fetch.total_s)); // every URL's download + for (slots.items, 0..) |sl, k| { + const url = sl.url orelse return cx.refuse(if (sl.kind == .image) "an image part needs an image_url with a url" else "a video part needs a video_url with a url"); + const bytes = try dataBytes(srv, cx, url, @min(if (sl.kind == .image) max_bytes_one else max_video_bytes, max_bytes_all - total_bytes), sl.kind, deadline); + total_bytes += bytes.len; + if (total_bytes > max_bytes_all) return cx.refuse("a request's images and videos exceed the 20 MiB limit"); + var text: std.ArrayList(u8) = .empty; + switch (sl.kind) { + .image => { + var img = qi.decode(cx.a, bytes, .{}) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.ImageTooLarge => return cx.refuse("image dimensions exceed the decoded pixel limit"), + error.NoDecoder => return cx.refuse("image inputs need the shared JPEG and PNG decoder, which this server does not have yet"), + else => return cx.refuse("image bytes are invalid or unsupported; use JPEG or PNG"), + }; + total_pixels += @as(u64, img.w) * img.h; + if (total_pixels > max_pixels_all) return cx.refuse("a request's images exceed the decoded pixel limit"); + const cap = if (std.mem.eql(u8, sl.detail, "low")) @min(per, 256) else per; + if (srv.info.vision_nemotron) { // Bolt: one tile of its own size, ... before the text + const p = qi.nemotron.prepare(cx.a, img) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the image's size is outside what the vision tower reads"), + }; + img.deinit(cx.a); + if (p.tokens() > per) return cx.fail(.request, "an image needs {d} visual tokens and this request's images have {d} each: send fewer images, or start the server with a larger --vision-image-tokens", .{ p.tokens(), per }); + try images.append(cx.a, .{ .patches = p.patches, .grid = .{ 1, @intCast(p.gh), @intCast(p.gw) } }); + try text.appendSlice(cx.a, ""); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, ""); + try text.appendSlice(cx.a, ""); + texts[k] = text.items; + continue; + } + if (srv.info.vision_glm) { // GLM: its canvas within [min(16, cap), cap] tokens, <|image|> placeholders + const p = qi.glm.prepare(cx.a, img, @min(16, cap), cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the image does not fit the visual-token budget"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + texts[k] = text.items; + continue; + } + const p = qi.prepare(cx.a, img, cap) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("absolute aspect ratio must be smaller than 200"), + }; + img.deinit(cx.a); + try images.append(cx.a, .{ .patches = p.patches, .grid = p.grid }); + try text.appendSlice(cx.a, vision_start); + for (0..p.tokens()) |_| try text.appendSlice(cx.a, "<|image_pad|>"); + try text.appendSlice(cx.a, vision_end); + }, + .video => if (srv.info.vision_nemotron) { + return cx.refuse("video inputs are not supported by the Nemotron CUDA engine yet: send images"); + } else if (srv.info.vision_glm) { + // GLM: <|begin_of_video|>, image groups with their "t.t seconds", <|end_of_video|> (glm5_next) + var v = try videoGlm(srv, cx, bytes); + defer v.deinit(cx.a); + try text.appendSlice(cx.a, "<|begin_of_video|>"); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + try text.appendSlice(cx.a, "<|begin_of_image|>"); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|image|>"); + try text.appendSlice(cx.a, "<|end_of_image|>"); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds"); + } + try text.appendSlice(cx.a, "<|end_of_video|>"); + } else { + var v = try video(srv, cx, bytes); + defer v.deinit(cx.a); + for (0..v.groups) |g| { + try images.append(cx.a, .{ .patches = try cx.a.dupe(f32, v.groupPatches(g)), .grid = .{ 1, @intCast(v.gh), @intCast(v.gw) } }); + var tb: [32]u8 = undefined; + try text.appendSlice(cx.a, "<"); + try text.appendSlice(cx.a, qi.formatSeconds(&tb, v.times[g]) catch unreachable); + try text.appendSlice(cx.a, " seconds>"); + try text.appendSlice(cx.a, vision_start); + for (0..v.tokens()) |_| try text.appendSlice(cx.a, "<|video_pad|>"); + try text.appendSlice(cx.a, vision_end); + } + }, + } + texts[k] = text.items; + } + // the messages with each media part as its text + const out = try cx.a.alloc(Value, list.array.len); + @memcpy(out, list.array); + var k: usize = 0; + while (k < slots.items.len) { + const i = slots.items[k].msg; + const c = list.array[i].get("content").?; + const parts = try cx.a.dupe(Value, c.array); + while (k < slots.items.len and slots.items[k].msg == i) : (k += 1) { + const t = try json.newObject(cx.a); + try t.put(cx.a, "type", .{ .string = "text" }); + try t.put(cx.a, "text", .{ .string = texts[k] }); + parts[slots.items[k].part] = .{ .object = t }; + } + const copy = try json.copyObject(cx.a, list.array[i].object); + try copy.put(cx.a, "content", .{ .array = parts }); + out[i] = .{ .object = copy }; + } + return .{ .messages = .{ .array = out }, .images = images.items }; +} + +// -- videos: libtfvideo (FFmpeg), loaded when the first video arrives ------------------------------------------ + +const TfVideo = extern struct { rgb: ?[*]u8, frames: c_int, height: c_int, width: c_int, indices: ?[*]c_int, fps: f64 }; +const SizeFn = *const fn (ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int; +const DecodeFn = *const fn (data: [*]const u8, len: usize, rate: f64, min_frames: c_int, max_frames: c_int, max_seconds: f64, max_dim: c_int, size: SizeFn, ctx: ?*anyopaque, out: *TfVideo) callconv(.c) c_int; +const FreeFn = *const fn (v: *TfVideo) callconv(.c) void; + +var lib_mutex: std.Io.Mutex = .init; +var lib: ?struct { decode: DecodeFn, free: FreeFn } = null; + +fn library(io: std.Io) ?@TypeOf(lib.?) { + lib_mutex.lockUncancelable(io); + defer lib_mutex.unlock(io); + if (lib == null) { + var dl = std.DynLib.open("libtfvideo.so") catch return null; + const d = dl.lookup(DecodeFn, "tf_video_decode") orelse return null; + const f = dl.lookup(FreeFn, "tf_video_free") orelse return null; + lib = .{ .decode = d, .free = f }; // kept open for the server's life + } + return lib; +} + +fn sizeFor(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const budget: *const u64 = @ptrCast(@alignCast(ctx.?)); + const s = qi.videoSize(@intCast(frames), @intCast(height), @intCast(width), budget.*) catch return -1; + oh.* = @intCast(s[0]); + ow.* = @intCast(s[1]); + return 0; +} + +/// GLM's video canvas: the frames' content size for libtfvideo (geometry within the budget), the canvas kept. +const GlmSize = struct { budget: u64, th: u32 = 0, tw: u32 = 0 }; + +fn sizeForGlm(ctx: ?*anyopaque, frames: c_int, height: c_int, width: c_int, oh: *c_int, ow: *c_int) callconv(.c) c_int { + const gs: *GlmSize = @ptrCast(@alignCast(ctx.?)); + const g = qi.glm.geometry(@intCast(frames), @intCast(height), @intCast(width), 16, gs.budget) catch return -1; + gs.th = g[0]; + gs.tw = g[1]; + oh.* = @intCast(g[2]); + ow.* = @intCast(g[3]); + return 0; +} + +fn videoGlm(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.glm.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var gs: GlmSize = .{ .budget = 16384 }; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| gs.budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch gs.budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 2, 256, 3600.0, 8192, sizeForGlm, @ptrCast(&gs), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.glm.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, gs.th, gs.tw, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} + +fn video(srv: *Server, cx: *Cx, bytes: []const u8) errors.Refused!qi.Video { + const l = library(srv.io) orelse return cx.refuse("video inputs need libtfvideo (FFmpeg) beside the server"); + var budget: u64 = 16384; + if (std.c.getenv("TENSORFOLD_VIDEO_TOKENS")) |v| budget = std.fmt.parseInt(u64, std.mem.span(v), 10) catch budget; + var v: TfVideo = undefined; + const rc = l.decode(bytes.ptr, bytes.len, 2.0, 4, 256, 3600.0, 8192, sizeFor, @ptrCast(&budget), &v); + if (rc != 0) return cx.refuse(switch (rc) { + -2 => "the video has no video stream", + -3 => "video dimensions are missing or exceed the pixel limit", + -4 => "the video's frame count or rate is missing", + -5 => "videos are limited to 60 minutes", + -6 => "the video has fewer than two decodable frames", + -8 => "a video needs at least 2 frames and an aspect ratio under 200", + else => "video bytes are invalid or unsupported; use MP4 or WebM", + }); + defer l.free(&v); + const frames: usize = @intCast(v.frames); + const h: usize = @intCast(v.height); + const w: usize = @intCast(v.width); + const idx = try cx.a.alloc(i32, frames); + for (idx, 0..) |*x, i| x.* = v.indices.?[i]; + return qi.prepareVideo(cx.a, v.rgb.?[0 .. frames * h * w * 3], frames, h, w, idx, v.fps) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + else => return cx.refuse("the video's frames do not fit the tower's patch grid"), + }; +} diff --git a/zig/src/vision/fixtures/nemotron_resize.bin b/zig/src/vision/fixtures/nemotron_resize.bin new file mode 100644 index 0000000000..d708a5eb5d Binary files /dev/null and b/zig/src/vision/fixtures/nemotron_resize.bin differ diff --git a/zig/src/vision/glm_image.zig b/zig/src/vision/glm_image.zig new file mode 100644 index 0000000000..889d3e679b --- /dev/null +++ b/zig/src/vision/glm_image.zig @@ -0,0 +1,190 @@ +//! GLM-5.3-Flash's image and video preparation as mlx-vlm's glm5_next/processing.py: canvas, resize, norm, patches. + +const std = @import("std"); +const qi = @import("qwen_image.zig"); + +pub const PATCH = 14; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; +const MEAN = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +const STD = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; +const PER = 3 * TEMPORAL * PATCH * PATCH; + +pub const Error = qi.Error; + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const f = @floor(x); + const d = x - f; + if (d > 0.5) return f + 1; + if (d < 0.5) return f; + return if (@mod(f, 2) == 0) f else f + 1; +} + +fn alignUp(v: f64, f: f64) f64 { + return @ceil(v / f) * f; +} + +/// smart_resize: the canvas (height, width) of ``frames`` frames within [min_tokens, max_tokens] merged tokens. +pub fn smartResize(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![2]u32 { + if (frames == 0 or height == 0 or width == 0 or min_tokens == 0 or min_tokens > max_tokens) return error.UnsupportedImage; + const f: f64 = FACTOR; + const ppt: f64 = TEMPORAL * f * f; + const min_px = @as(f64, @floatFromInt(min_tokens)) * ppt; + const max_px = @as(f64, @floatFromInt(max_tokens)) * ppt; + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const n: f64 = @floatFromInt(frames); + const af = @max(@as(f64, TEMPORAL), pyRound(n / TEMPORAL) * TEMPORAL); + var ah = alignUp(h, f); + var aw = alignUp(w, f); + var px = af * ah * aw; + if (px < min_px) { + const scale = @sqrt(min_px / (n * h * w)); + ah = alignUp(@max(1, @ceil(h * scale)), f); + aw = alignUp(@max(1, @ceil(w * scale)), f); + px = af * ah * aw; + } + if (px > max_px) { // the tallest content within the budget, by bisection on its height + if (max_px < af * f * f) return error.UnsupportedImage; + var low: f64 = 1; + var high: f64 = h; + var best: [2]f64 = .{ f, f }; + while (low <= high) { + const ch = @floor((low + high) / 2); + const cw = @max(1, @floor(w * ch / h)); + const cand = [2]f64{ alignUp(ch, f), alignUp(cw, f) }; + if (af * cand[0] * cand[1] <= max_px) { + best = cand; + low = ch + 1; + } else high = ch - 1; + } + ah = best[0]; + aw = best[1]; + } + return .{ @intFromFloat(ah), @intFromFloat(aw) }; +} + +/// _resize_geometry: the canvas (th, tw) and the content (ch, cw) resized into its top-left corner. +pub fn geometry(frames: u32, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![4]u32 { + const t = try smartResize(frames, height, width, min_tokens, max_tokens); + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + var scale = @min(@as(f64, @floatFromInt(t[0])) / h, @as(f64, @floatFromInt(t[1])) / w); + const ppt: f64 = TEMPORAL * FACTOR * FACTOR; + if (@as(f64, @floatFromInt(frames)) * h * w >= @as(f64, @floatFromInt(min_tokens)) * ppt) scale = @min(1.0, scale); + const ch: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[0])), @floor(h * scale)))); + const cw: u32 = @intFromFloat(@max(1, @min(@as(f64, @floatFromInt(t[1])), @floor(w * scale)))); + return .{ t[0], t[1], ch, cw }; +} + +pub const Prepared = struct { + patches: []f32, // [gh * gw, 1176] + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], MERGE * MERGE)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +fn norm(v: u8, c: usize) f32 { + const x = @as(f32, @floatFromInt(v)) * @as(f32, @floatCast(@as(f64, 1.0) / 255.0)); + return (x - MEAN[c]) / STD[c]; +} + +/// A frame's normalized value at canvas (y, x) (zero padding past the content's ch x cw). +fn at(frame: []const u8, cw: usize, ch: usize, y: usize, x: usize, c: usize) f32 { + return norm(if (y < ch and x < cw) frame[(y * cw + x) * 3 + c] else 0, c); +} + +/// ``im`` on its canvas within [min_tokens, max_tokens] tokens, normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: qi.Image, min_tokens: u64, max_tokens: u64) Error!Prepared { + const g = try geometry(TEMPORAL, im.h, im.w, min_tokens, max_tokens); + var r = if (g[2] == im.h and g[3] == im.w) null else try qi.resize(gpa, im, g[3], g[2]); + defer if (r) |*x| x.deinit(gpa); + const src = if (r) |x| x.rgb else im.rgb; + const gh: usize = g[0] / PATCH; + const gw: usize = g[1] / PATCH; + const out = try gpa.alloc(f32, gh * gw * PER); + var n: usize = 0; + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..PATCH) |y| for (0..PATCH) |x| { + const v = at(src, g[3], g[2], py + y, px + x, c); + for (0..TEMPORAL) |t| dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = v; + }; + n += 1; + }; + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +/// A video's frame groups: patches of two consecutive frames (an odd last repeats) and the first frame's seconds. +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1176 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * PER; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames at the content size (ch x cw: geometry's), placed on the th x tw canvas. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, ch: usize, cw: usize, th: usize, tw: usize, indices: []const i32, fps: f64) Error!Video { + if (th % FACTOR != 0 or tw % FACTOR != 0 or ch > th or cw > tw or frames < 1) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = th / PATCH; + const gw = tw / PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * PER); + errdefer gpa.free(out); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * PER ..][0..PER]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * ch * cw * 3 ..][0 .. ch * cw * 3]; + for (0..PATCH) |y| for (0..PATCH) |x| { + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = at(frame, cw, ch, py + y, px + x, c); + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| times[g] = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +test "smart_resize and geometry match the processor's" { + // 640 x 480 within 8000 tokens: aligned to 504 x 644, content unscaled + try std.testing.expectEqual([4]u32{ 504, 644, 480, 640 }, try geometry(TEMPORAL, 480, 640, 16, 8000)); + // 40 x 30 below 16 tokens: scaled up, content scaled + const g = try geometry(TEMPORAL, 30, 40, 16, 8000); + try std.testing.expectEqual(@as(u32, 112), g[0]); + try std.testing.expectEqual(@as(u32, 140), g[1]); + // 300 x 900 (tall) + const t = try geometry(TEMPORAL, 900, 300, 16, 8000); + try std.testing.expectEqual([2]u32{ 924, 308 }, .{ t[0], t[1] }); +} diff --git a/zig/src/vision/nemotron_image.zig b/zig/src/vision/nemotron_image.zig new file mode 100644 index 0000000000..0abac37bc2 --- /dev/null +++ b/zig/src/vision/nemotron_image.zig @@ -0,0 +1,273 @@ +//! Bolt's image prep as vLLM runs it: one tile, torch's fp32 bicubic antialiased resize, CLIP norm, 16x16 patches. + +const std = @import("std"); +const qi = @import("qwen_image.zig"); + +pub const Error = qi.Error; +pub const Image = qi.Image; + +pub const PATCH = 16; +pub const MERGE = 2; // pixel shuffle v2: 2x2 patches a token (downsample_ratio 0.5) +/// One patch's values, (channel, row, column) order: patch_generator.embedder's input width. +pub const PER = 3 * PATCH * PATCH; +pub const MIN_PATCHES = 1024; +pub const MAX_PATCHES = 13312; +/// vLLM's budget, (max_model_len - 4) * 4 patches, before the max_num_patches cap (which always binds). +pub const MODEL_LEN_BUDGET = (16384 - 4) * MERGE * MERGE; +pub const MEAN = [3]f32{ 0.48145466, 0.4578275, 0.40821073 }; +pub const STD = [3]f32{ 0.26862954, 0.26130258, 0.27577711 }; + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const f = @floor(x); + const d = x - f; + if (d > 0.5) return f + 1; + if (d < 0.5) return f; + return if (@mod(f, 2) == 0) f else f + 1; +} + +/// The patch grid (rows, columns) of the single tile: target_patches in the reference, both sides even. +pub fn grid(width: u32, height: u32) Error![2]u32 { + if (width == 0 or height == 0) return error.UnsupportedImage; + const r: f64 = MERGE; + const mn: f64 = MIN_PATCHES; + const avail: f64 = @max(@min(@as(f64, MODEL_LEN_BUDGET), MAX_PATCHES), mn); + const ph = pyRound(@as(f64, @floatFromInt(height)) / PATCH + 0.5); + const pw = pyRound(@as(f64, @floatFromInt(width)) / PATCH + 0.5); + const f = @min(@sqrt(avail / (ph * pw)), 1.0); + var th = @floor(f * ph); + var tw = @floor(f * pw); + if (avail > mn and th * tw < mn) { + if (th * tw == 0) return error.BadAspectRatio; // Python divides by zero here + const up = @sqrt(mn / (th * tw)); + th = @ceil(up * th); + tw = @ceil(up * tw); + } + if (@mod(th, r) != 0) { + const up = th + (r - @mod(th, r)); + th = if (up * tw <= avail) up else @max(r, th - @mod(th, r)); + } + if (@mod(tw, r) != 0) { + const up = tw + (r - @mod(tw, r)); + tw = if (th * up <= avail) up else @max(r, tw - @mod(tw, r)); + } + return .{ @intFromFloat(th), @intFromFloat(tw) }; +} + +/// torch's antialiased bicubic filter (Keys, a = -0.5), in fp32 as aten's aa_filter computes it. +fn aaCubic(x0: f32) f32 { + const a: f32 = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0; + if (x < 2.0) return ((a * x - 5.0 * a) * x + 8.0 * a) * x - 4.0 * a; + return 0.0; +} + +/// One axis' taps: output i reads input [first[i], first[i] + count[i]) with weights w[i * ksize ..]. +const Taps = struct { + ksize: usize, + first: []usize, + count: []usize, + w: []f32, + + fn deinit(t: *Taps, gpa: std.mem.Allocator) void { + gpa.free(t.first); + gpa.free(t.count); + gpa.free(t.w); + } +}; + +/// aten's _compute_weights_aa for an axis (align_corners=False, no scale factor): fp32 throughout, then normalized. +fn taps(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Taps { + const scale: f32 = @as(f32, @floatFromInt(in_size)) / @as(f32, @floatFromInt(out_size)); + const support: f32 = if (scale >= 1.0) 2.0 * scale else 2.0; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const invscale: f32 = if (scale >= 1.0) 1.0 / scale else 1.0; + var t: Taps = .{ .ksize = ksize, .first = try gpa.alloc(usize, out_size), .count = undefined, .w = undefined }; + errdefer gpa.free(t.first); + t.count = try gpa.alloc(usize, out_size); + errdefer gpa.free(t.count); + t.w = try gpa.alloc(f32, out_size * ksize); + @memset(t.w, 0); + for (0..out_size) |i| { + const center: f32 = scale * (@as(f32, @floatFromInt(i)) + 0.5); + const lo: i64 = @max(@as(i64, @intFromFloat(center - support + 0.5)), 0); + const hi: i64 = @min(@as(i64, @intFromFloat(center + support + 0.5)), @as(i64, @intCast(in_size))); + const n: usize = @intCast(hi - lo); + const w = t.w[i * ksize ..][0..n]; + var total: f32 = 0; + for (w, 0..) |*v, j| { + const x: f32 = (@as(f32, @floatFromInt(@as(i64, @intCast(j)) + lo)) - center + 0.5) * invscale; + v.* = aaCubic(x); + total += v.*; + } + if (total != 0) for (w) |*v| { + v.* /= total; + }; + t.first[i] = @intCast(lo); + t.count[i] = n; + } + return t; +} + +/// Bicubic antialiased F.interpolate (align_corners=False) of fp32 [3, h, w] to [3, oh, ow]: width pass, then height. +pub fn resizePlanes(gpa: std.mem.Allocator, src: []const f32, w: u32, h: u32, ow: u32, oh: u32) Error![]f32 { + var cur = try gpa.dupe(f32, src); + var cw: usize = w; + if (ow != w) { + var t = try taps(gpa, w, ow); + defer t.deinit(gpa); + const out = try gpa.alloc(f32, 3 * @as(usize, h) * ow); + for (0..3 * @as(usize, h)) |row| { + const in = cur[row * cw ..][0..cw]; + for (0..ow) |x| { + const k = t.w[x * t.ksize ..]; + const s = in[t.first[x]..]; + var acc: f32 = s[0] * k[0]; + for (1..t.count[x]) |j| acc += s[j] * k[j]; + out[row * ow + x] = acc; + } + } + gpa.free(cur); + cur = out; + cw = ow; + } + if (oh != h) { + var t = try taps(gpa, h, oh); + defer t.deinit(gpa); + const out = try gpa.alloc(f32, 3 * @as(usize, oh) * cw); + for (0..3) |c| { + const plane = cur[c * h * cw ..][0 .. h * cw]; + for (0..oh) |y| { + const k = t.w[y * t.ksize ..]; + const o = out[(c * oh + y) * cw ..][0..cw]; + for (o, 0..) |*v, x| { + var acc: f32 = plane[t.first[y] * cw + x] * k[0]; + for (1..t.count[y]) |j| acc += plane[(t.first[y] + j) * cw + x] * k[j]; + v.* = acc; + } + } + } + gpa.free(cur); + cur = out; + } + return cur; +} + +/// RGB bytes resized to [3, oh, ow] fp32 and normalized as the reference does: (x / 255 - mean) / std in fp32. +pub fn pixelValues(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error![]f32 { + const n = @as(usize, im.w) * im.h; + if (im.rgb.len != n * 3) return error.UnsupportedImage; + const planes = try gpa.alloc(f32, 3 * n); + defer gpa.free(planes); + for (0..n) |i| for (0..3) |c| { + planes[c * n + i] = @floatFromInt(im.rgb[i * 3 + c]); + }; + const out = try resizePlanes(gpa, planes, im.w, im.h, ow, oh); + for (0..3) |c| for (out[c * @as(usize, ow) * oh ..][0 .. @as(usize, ow) * oh]) |*v| { + v.* = (v.* / 255.0 - MEAN[c]) / STD[c]; + }; + return out; +} + +/// One image ready for the tower: its patches (fp32 [gh * gw, 768], row-major over the grid, each patch (c, dy, dx)). +pub const Prepared = struct { + patches: []f32, + gh: u32, + gw: u32, + + pub fn count(p: Prepared) usize { + return @as(usize, p.gh) * p.gw; + } + + /// The `` placeholders the prompt carries: one a 2x2 group of patches. + pub fn tokens(p: Prepared) usize { + return p.count() / (MERGE * MERGE); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +/// Planes [3, gh * 16, gw * 16] cut into patches, (c, dy, dx) within each (the reference's reshape and permute). +pub fn patchify(gpa: std.mem.Allocator, px: []const f32, gh: u32, gw: u32) Error![]f32 { + const H = @as(usize, gh) * PATCH; + const W = @as(usize, gw) * PATCH; + if (px.len != 3 * H * W) return error.UnsupportedImage; + const out = try gpa.alloc(f32, @as(usize, gh) * gw * PER); + for (0..gh) |py| for (0..gw) |pxi| { + const dst = out[(py * gw + pxi) * PER ..][0..PER]; + for (0..3) |c| for (0..PATCH) |dy| { + @memcpy(dst[(c * PATCH + dy) * PATCH ..][0..PATCH], px[(c * H + py * PATCH + dy) * W + pxi * PATCH ..][0..PATCH]); + }; + }; + return out; +} + +/// An RGB image to the tower's patches: grid, resize, normalize, patchify. +pub fn prepare(gpa: std.mem.Allocator, im: Image) Error!Prepared { + const g = try grid(im.w, im.h); + const px = try pixelValues(gpa, im, g[1] * PATCH, g[0] * PATCH); + defer gpa.free(px); + return .{ .patches = try patchify(gpa, px, g[0], g[1]), .gh = g[0], .gw = g[1] }; +} + +const fixture = @embedFile("fixtures/nemotron_resize.bin"); + +fn testImage(gpa: std.mem.Allocator, w: u32, h: u32) !Image { + const rgb = try gpa.alloc(u8, @as(usize, w) * h * 3); + for (0..h) |y| for (0..w) |x| for (0..3) |c| { + rgb[(y * w + x) * 3 + c] = @intCast((7 * x + 13 * y + 29 * c + (x * y) % 17) % 256); + }; + return .{ .w = w, .h = h, .rgb = rgb }; +} + +test "the patch grid matches target_patches (tools/zig/bolt_image_fixture.py's table)" { + const cases = [_][4]u32{ + .{ 2172, 724, 136, 46 }, .{ 640, 480, 40, 30 }, .{ 37, 23, 40, 28 }, .{ 4000, 3000, 132, 100 }, + .{ 16, 16, 32, 32 }, .{ 100, 5000, 8, 314 }, .{ 1024, 1024, 64, 64 }, .{ 512, 512, 32, 32 }, + .{ 3000, 100, 188, 8 }, .{ 1920, 1080, 120, 68 }, .{ 800, 600, 50, 38 }, + }; + for (cases) |c| { + const g = try grid(c[0], c[1]); + try std.testing.expectEqual(c[3], g[0]); + try std.testing.expectEqual(c[2], g[1]); + try std.testing.expect(@as(usize, g[0]) * g[1] <= MAX_PATCHES); + } +} + +test "resize and normalization match torch's bicubic antialias (up and down) to 5e-6" { + const gpa = std.testing.allocator; + var im = try testImage(gpa, 37, 23); + defer im.deinit(gpa); + const want: []align(1) const f32 = std.mem.bytesAsSlice(f32, fixture); + var at: usize = 0; + var worst: f32 = 0; + for ([_][2]u32{ .{ 48, 32 }, .{ 16, 16 }, .{ 24, 8 } }) |s| { + const got = try pixelValues(gpa, im, s[0], s[1]); + defer gpa.free(got); + for (got, want[at..][0..got.len]) |g, w| worst = @max(worst, @abs(g - w)); + at += got.len; + } + try std.testing.expectEqual(want.len, at); + try std.testing.expect(worst < 5e-6); +} + +test "patches are (c, dy, dx) over a row-major grid; prepare's count and tokens" { + const gpa = std.testing.allocator; + const px = try gpa.alloc(f32, 3 * 32 * 48); + defer gpa.free(px); + for (px, 0..) |*v, i| v.* = @floatFromInt(i); + const p = try patchify(gpa, px, 2, 3); + defer gpa.free(p); + // patch (1, 2), channel 2, dy 5, dx 7 is pixel (2, 16 + 5, 32 + 7) + try std.testing.expectEqual(@as(f32, @floatFromInt((2 * 32 + 21) * 48 + 39)), p[(1 * 3 + 2) * PER + (2 * PATCH + 5) * PATCH + 7]); + var im = try testImage(gpa, 37, 23); + defer im.deinit(gpa); + var prep = try prepare(gpa, im); + defer prep.deinit(gpa); + try std.testing.expectEqual(@as(usize, 28 * 40), prep.count()); + try std.testing.expectEqual(@as(usize, 280), prep.tokens()); + try std.testing.expectEqual(prep.count() * PER, prep.patches.len); +} diff --git a/zig/src/vision/qwen_image.zig b/zig/src/vision/qwen_image.zig new file mode 100644 index 0000000000..774a22ccac --- /dev/null +++ b/zig/src/vision/qwen_image.zig @@ -0,0 +1,364 @@ +//! Image inputs for Qwen-style vision towers as Qwen2VL's processor makes them: decode, smart_resize, resize, patches. + +const std = @import("std"); +/// GLM-5.3-Flash's preparation (its canvas, padding, normalization and patch order) on this module's decode and resize. +pub const glm = @import("glm_image.zig"); +/// Nemotron 3.5 Bolt's preparation (one dynamic-resolution tile, torch's fp32 resize, C-RADIO's patches). +pub const nemotron = @import("nemotron_image.zig"); + +pub const Limits = struct { + max_dimension: u32 = 8192, + max_pixels: u64 = 16 * 1024 * 1024, +}; + +/// An RGB image, rows top to bottom, three bytes a pixel. +pub const Image = struct { + w: u32, + h: u32, + rgb: []u8, + + pub fn deinit(im: *Image, gpa: std.mem.Allocator) void { + gpa.free(im.rgb); + } +}; + +pub const Error = error{ UnsupportedImage, ImageTooLarge, OutOfMemory, BadAspectRatio, NoDecoder }; + +/// JPEG or PNG bytes to RGB: the shared runtime decoder's job (pinned libjpeg-turbo); until it lands, refused by name. +pub fn decode(gpa: std.mem.Allocator, bytes: []const u8, limits: Limits) Error!Image { + _ = .{ gpa, bytes, limits }; + return error.NoDecoder; +} + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + const r = @round(x); + if (@abs(x - @trunc(x)) == 0.5) return 2.0 * @round(x / 2.0); + return r; +} + +/// transformers' smart_resize: both sides multiples of ``factor``, the area within [min_pixels, max_pixels]. +pub fn smartResize(height: u32, width: u32, factor: u32, min_pixels: u64, max_pixels: u64) Error![2]u32 { + const h: f64 = @floatFromInt(height); + const w: f64 = @floatFromInt(width); + const f: f64 = @floatFromInt(factor); + if (@max(h, w) / @min(h, w) > 200) return error.BadAspectRatio; + var hb = pyRound(h / f) * f; + var wb = pyRound(w / f) * f; + if (hb * wb > @as(f64, @floatFromInt(max_pixels))) { + const beta = @sqrt(h * w / @as(f64, @floatFromInt(max_pixels))); + hb = @max(f, @floor(h / beta / f) * f); + wb = @max(f, @floor(w / beta / f) * f); + } else if (hb * wb < @as(f64, @floatFromInt(min_pixels))) { + const beta = @sqrt(@as(f64, @floatFromInt(min_pixels)) / (h * w)); + hb = @ceil(h * beta / f) * f; + wb = @ceil(w * beta / f) * f; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + +fn bicubic(x0: f64) f64 { + const a = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1; + if (x < 2.0) return (((x - 5) * x + 8) * x - 4) * a; + return 0.0; +} + +const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; + +/// PIL's precompute_coeffs for an axis, as int16 fixed point at the most precision (< 22 bits) torch's uint8 takes. +fn coeffs(gpa: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { + const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); + const filterscale = @max(scale, 1.0); + const support = 2.0 * filterscale; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const bounds = try gpa.alloc(usize, out_size * 2); + const kk = try gpa.alloc(i32, out_size * ksize); + const pre_all = try gpa.alloc(f64, out_size * ksize); + defer gpa.free(pre_all); + for (0..out_size) |xx| { + const pre = pre_all[xx * ksize ..][0..ksize]; + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + var ww: f64 = 0; + const ss = 1.0 / filterscale; + var xmin_i: i64 = @intFromFloat(center - support + 0.5); + if (xmin_i < 0) xmin_i = 0; + var xmax_i: i64 = @intFromFloat(center + support + 0.5); + if (xmax_i > @as(i64, @intCast(in_size))) xmax_i = @intCast(in_size); + const xmin: usize = @intCast(xmin_i); + const xmax: usize = @intCast(xmax_i - xmin_i); + @memset(pre, 0); + for (0..xmax) |x| { + const wgt = bicubic((@as(f64, @floatFromInt(x + xmin)) - center + 0.5) * ss); + pre[x] = wgt; + ww += wgt; + } + for (0..xmax) |x| { + if (ww != 0.0) pre[x] /= ww; + } + bounds[xx * 2] = xmin; + bounds[xx * 2 + 1] = xmax; + } + var w_max: f64 = 0; + for (pre_all) |v| w_max = @max(w_max, @abs(v)); + var precision: u5 = 0; + while (precision < 22) : (precision += 1) { + const next: i64 = @intFromFloat(0.5 + w_max * @as(f64, @floatFromInt(@as(i64, 1) << (precision + 1)))); + if (next >= (1 << 15)) break; + } + const one: f64 = @floatFromInt(@as(i64, 1) << precision); + for (pre_all, 0..) |v, i| kk[i] = @intFromFloat(if (v < 0) -0.5 + v * one else 0.5 + v * one); + return .{ .ksize = ksize, .bounds = bounds, .kk = kk, .precision = precision }; +} + +fn clip8(v: i64, precision: u5) u8 { + const r = v >> precision; + return @intCast(std.math.clamp(r, 0, 255)); +} + +/// The antialiased bicubic resize (PIL's ImagingResample, 8 bits): horizontal pass, then vertical, rounding to bytes. +pub fn resize(gpa: std.mem.Allocator, im: Image, ow: u32, oh: u32) Error!Image { + var cur = im; + var owned = false; + defer if (owned) cur.deinit(gpa); + if (ow != im.w) { + const c = try coeffs(gpa, im.w, ow); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, ow) * im.h * 3); + for (0..im.h) |y| for (0..ow) |x| { + const xmin = c.bounds[x * 2]; + const xmax = c.bounds[x * 2 + 1]; + const k = c.kk[x * c.ksize ..]; + for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..xmax) |j| ss += @as(i64, cur.rgb[(y * cur.w + xmin + j) * 3 + ch]) * k[j]; + out[(y * ow + x) * 3 + ch] = clip8(ss, c.precision); + } + }; + cur = .{ .w = ow, .h = im.h, .rgb = out }; + owned = true; + } + if (oh != cur.h) { + const c = try coeffs(gpa, cur.h, oh); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, cur.w) * oh * 3); + for (0..oh) |y| { + const ymin = c.bounds[y * 2]; + const ymax = c.bounds[y * 2 + 1]; + const k = c.kk[y * c.ksize ..]; + for (0..cur.w) |x| for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..ymax) |j| ss += @as(i64, cur.rgb[((ymin + j) * cur.w + x) * 3 + ch]) * k[j]; + out[(y * cur.w + x) * 3 + ch] = clip8(ss, c.precision); + }; + } + if (owned) cur.deinit(gpa); + cur = .{ .w = cur.w, .h = oh, .rgb = out }; + owned = true; + } + if (!owned) return .{ .w = im.w, .h = im.h, .rgb = try gpa.dupe(u8, im.rgb) }; + owned = false; + return cur; +} + +/// One image prepared for the tower: its patches (fp32 [gh * gw, 1536], merge-block order) and grid [1, gh, gw]. +pub const Prepared = struct { + patches: []f32, + grid: [3]i64, + + pub fn tokens(p: Prepared) usize { + return @intCast(@divExact(p.grid[1] * p.grid[2], 4)); + } + + pub fn deinit(p: *Prepared, gpa: std.mem.Allocator) void { + gpa.free(p.patches); + } +}; + +pub const PATCH = 16; +pub const MERGE = 2; +pub const TEMPORAL = 2; +pub const FACTOR = PATCH * MERGE; + +/// ``im`` resized within ``max_tokens`` visual tokens (min_pixels: one merged patch), normalized and cut into patches. +pub fn prepare(gpa: std.mem.Allocator, im: Image, max_tokens: u64) Error!Prepared { + const cap = max_tokens * FACTOR * FACTOR; + const size = try smartResize(im.h, im.w, FACTOR, @min(FACTOR * FACTOR, cap), cap); + var r = try resize(gpa, im, size[1], size[0]); + defer r.deinit(gpa); + const gh: usize = size[0] / PATCH; + const gw: usize = size[1] / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, gh * gw * per); + var n: usize = 0; + var bh: usize = 0; + while (bh < gh / MERGE) : (bh += 1) { + var bw: usize = 0; + while (bw < gw / MERGE) : (bw += 1) { + for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @floatFromInt(r.rgb[((py + y) * r.w + px + x) * 3 + c]); + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 127.5) / 127.5; + }; + n += 1; + }; + } + } + return .{ .patches = out, .grid = .{ 1, @intCast(gh), @intCast(gw) } }; +} + +// -- videos --------------------------------------------------------------------------------------------------------- + +/// A video's frame size: Qwen3-VL's smart_resize, at most 768 tokens a frame group, the video within ``budget``. +pub fn videoSize(frames: u32, height0: u32, width0: u32, budget: u64) Error![2]u32 { + const factor: f64 = FACTOR; + const shortest: f64 = 128 * FACTOR * FACTOR; + const longest: f64 = @as(f64, @floatFromInt(budget)) * TEMPORAL * FACTOR * FACTOR; + const per_frame = @max(@min(768 * FACTOR * FACTOR, @floor(longest / @as(f64, @floatFromInt(@max(1, frames))))), @floor(shortest * 1.05)); + const max_pixels = per_frame * @as(f64, @floatFromInt(frames)); + if (frames < TEMPORAL) return error.UnsupportedImage; + var height: f64 = @floatFromInt(height0); + var width: f64 = @floatFromInt(width0); + if (height < factor or width < factor) { + const scale = @max(factor / height, factor / width); + height = @trunc(height * scale); + width = @trunc(width * scale); + } + if (@max(height, width) / @min(height, width) > 200) return error.BadAspectRatio; + var hb = pyRound(height / factor) * factor; + var wb = pyRound(width / factor) * factor; + const tb = pyRound(@as(f64, @floatFromInt(frames)) / TEMPORAL) * TEMPORAL; + const n: f64 = @floatFromInt(frames); + if (tb * hb * wb > max_pixels) { + const beta = @sqrt(n * height * width / max_pixels); + hb = @max(factor, @floor(height / beta / factor) * factor); + wb = @max(factor, @floor(width / beta / factor) * factor); + } else if (tb * hb * wb < shortest) { + const beta = @sqrt(shortest / (n * height * width)); + hb = @ceil(height * beta / factor) * factor; + wb = @ceil(width * beta / factor) * factor; + } + return .{ @intFromFloat(hb), @intFromFloat(wb) }; +} + +/// A video's frame groups: patches of two consecutive frames, normalized as numpy does, and their mean time. +pub const Video = struct { + patches: []f32, // groups x gh x gw x 1536 + groups: usize, + gh: usize, + gw: usize, + times: []f64, + + pub fn groupPatches(v: Video, g: usize) []f32 { + const per = v.gh * v.gw * 3 * TEMPORAL * PATCH * PATCH; + return v.patches[g * per ..][0..per]; + } + + pub fn tokens(v: Video) usize { + return v.gh * v.gw / (MERGE * MERGE); + } + + pub fn deinit(v: *Video, gpa: std.mem.Allocator) void { + gpa.free(v.patches); + gpa.free(v.times); + } +}; + +/// ``rgb``: ``frames`` RGB frames of ``h`` x ``w`` (already at the tower's size), their source ``indices`` and ``fps``. +pub fn prepareVideo(gpa: std.mem.Allocator, rgb: []const u8, frames: usize, h: usize, w: usize, indices: []const i32, fps: f64) Error!Video { + if (h % FACTOR != 0 or w % FACTOR != 0 or frames < 2) return error.UnsupportedImage; + const groups = (frames + TEMPORAL - 1) / TEMPORAL; + const gh = h / PATCH; + const gw = w / PATCH; + const per = 3 * TEMPORAL * PATCH * PATCH; + const out = try gpa.alloc(f32, groups * gh * gw * per); + errdefer gpa.free(out); + const scale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var n: usize = 0; + for (0..groups) |g| { + for (0..gh / MERGE) |bh| for (0..gw / MERGE) |bw| for (0..MERGE) |mh| for (0..MERGE) |mw| { + const py = (bh * MERGE + mh) * PATCH; + const px = (bw * MERGE + mw) * PATCH; + const dst = out[n * per ..][0..per]; + for (0..3) |c| for (0..TEMPORAL) |t| { + const f = @min(g * TEMPORAL + t, frames - 1); + const frame = rgb[f * h * w * 3 ..]; + for (0..PATCH) |y| for (0..PATCH) |x| { + const v: f32 = @as(f32, @floatFromInt(frame[((py + y) * w + px + x) * 3 + c])) * scale; + dst[((c * TEMPORAL + t) * PATCH + y) * PATCH + x] = (v - 0.5) / 0.5; + }; + }; + n += 1; + }; + } + const times = try gpa.alloc(f64, groups); + for (0..groups) |g| { + const a: f64 = @as(f64, @floatFromInt(indices[g * TEMPORAL])) / fps; + const b: f64 = @as(f64, @floatFromInt(indices[@min(g * TEMPORAL + TEMPORAL - 1, frames - 1)])) / fps; + times[g] = (a + b) / 2; + } + return .{ .patches = out, .groups = groups, .gh = gh, .gw = gw, .times = times }; +} + +/// Python's f"{t:.1f}": one decimal, an exact tie rounded to even. +pub fn formatSeconds(buf: []u8, t: f64) ![]const u8 { + const v = tenths(t); + return std.fmt.bufPrint(buf, "{d}.{d}", .{ @divTrunc(v, 10), @as(u64, @intCast(@mod(v, 10))) }); +} + +/// t times ten, an exact tie to even, in integers: t is m * 2^e and 10m fits 57 bits (f128 would link compiler-rt). +fn tenths(t: f64) i64 { + const bits: u64 = @bitCast(@abs(t)); + const biased: i32 = @intCast(bits >> 52); + const m: u64 = if (biased == 0) bits & ((1 << 52) - 1) else (bits & ((1 << 52) - 1)) | (1 << 52); + const e = @as(i32, @max(biased, 1)) - 1075; + if (e >= 0) return @intFromFloat(t * 10); // at least 2^52 seconds: never a frame's time + const tm = m * 10; + if (e <= -58) return 0; // t * 10 under a half: tm < 2^57 <= 2^(s - 1) + const s: u6 = @intCast(-e); + const q = tm >> s; + const rem = tm & ((@as(u64, 1) << s) - 1); + const half = @as(u64, 1) << (s - 1); + const r: i64 = @intCast(if (rem > half or (rem == half and q & 1 == 1)) q + 1 else q); + return if (t < 0) -r else r; +} + +test "tenths rounds t * 10 exactly, ties to even, as the f128 product does" { + const ref = struct { + fn f(t: f64) i64 { + const x: f128 = @as(f128, t) * 10; + const fl = @floor(x); + const frac = x - fl; + var r = if (frac > 0.5) fl + 1 else fl; + if (frac == 0.5 and @rem(fl, 2) != 0) r = fl + 1; + return @intFromFloat(r); + } + }.f; + const fixed = [_]f64{ 0, -0.0, 0.05, 0.15, 0.25, 0.35, 0.45, 1.25, 2.5, 0.95, 1.05, 3.0, 123.45, 1e-300, 5e-324, 4503599627370495.5, -0.25, -1.35 }; + for (fixed) |t| try std.testing.expectEqual(ref(t), tenths(t)); + var prng: std.Random.DefaultPrng = .init(0x7e17); + const rand = prng.random(); + for (0..200_000) |i| { + const t: f64 = switch (i % 4) { + 0 => rand.float(f64) * 4000, + 1 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 0, 80_000))) / 20, // halves of tenths + 2 => @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 100_000))) / @as(f64, @floatFromInt(rand.intRangeLessThan(u32, 1, 121))), + else => @bitCast(@as(u64, rand.uintLessThan(u64, 1075)) << 52 | rand.int(u64) >> 12), // any double below 2^52 + }; + try std.testing.expectEqual(ref(t), tenths(t)); + } + var buf: [32]u8 = undefined; + try std.testing.expectEqualStrings("0.2", try formatSeconds(&buf, 0.25)); + try std.testing.expectEqualStrings("1.4", try formatSeconds(&buf, 1.35)); +} + +test { + _ = glm; + _ = nemotron; +} diff --git a/zig/src/vision/tf_video.c b/zig/src/vision/tf_video.c new file mode 100644 index 0000000000..9cd1290cc7 --- /dev/null +++ b/zig/src/vision/tf_video.c @@ -0,0 +1,211 @@ +// libtfvideo: a video's sampled frames as videos.py decodes them with PyAV 19.0.1 (FFmpeg 9), RGB via SWS_BICUBIC. +#include +#include +#include +#include +#include +#include +#include +#include + +typedef int (*tf_size_fn)(void* ctx, int frames, int height, int width, int* out_h, int* out_w); + +typedef struct { + uint8_t* rgb; // frames x height x width x 3 + int frames, height, width; + int* indices; // each frame's index in the source + double fps; // the source's rate +} tf_video; + +enum { + TF_OK = 0, + TF_INVALID = -1, // not a decodable container or stream + TF_NO_STREAM = -2, // no video stream + TF_DIMENSIONS = -3, // missing or over max_dim + TF_NO_COUNT = -4, // the frame count or rate is missing + TF_TOO_LONG = -5, // over max_seconds + TF_FEW_FRAMES = -6, // fewer than two decodable frames + TF_MEMORY = -7, + TF_SIZE = -8, // the size callback refused +}; + +typedef struct { const uint8_t* data; size_t len, pos; } mem_t; + +static int mem_read(void* o, uint8_t* buf, int n) { + mem_t* m = (mem_t*)o; + if (m->pos >= m->len) return AVERROR_EOF; + size_t k = m->len - m->pos < (size_t)n ? m->len - m->pos : (size_t)n; + memcpy(buf, m->data + m->pos, k); + m->pos += k; + return (int)k; +} + +static int64_t mem_seek(void* o, int64_t off, int whence) { + mem_t* m = (mem_t*)o; + if (whence == AVSEEK_SIZE) return (int64_t)m->len; + int64_t at = whence == SEEK_SET ? off : whence == SEEK_CUR ? (int64_t)m->pos + off : (int64_t)m->len + off; + if (at < 0 || at > (int64_t)m->len) return -1; + m->pos = (size_t)at; + return at; +} + +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): the frame's own AVColorSpace passes through. +static void set_colorspace(AVFrame* f, int cs) { + if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; + else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; + else if (cs == SWS_CS_ITU601) f->colorspace = AVCOL_SPC_SMPTE170M; + else if (cs == SWS_CS_SMPTE240M) f->colorspace = AVCOL_SPC_SMPTE240M; + else if (cs == SWS_CS_BT2020) f->colorspace = AVCOL_SPC_BT2020_NCL; +} + +static double round_half_even(double x) { + double r = round(x); + if (fabs(x - trunc(x)) == 0.5) r = 2.0 * round(x / 2.0); + return r; +} + +void tf_video_free(tf_video* v) { + if (!v) return; + free(v->rgb); + free(v->indices); + v->rgb = NULL; + v->indices = NULL; +} + +int tf_video_decode(const uint8_t* data, size_t len, double rate, int min_frames, int max_frames, double max_seconds, + int max_dim, tf_size_fn size, void* ctx, tf_video* out) { + memset(out, 0, sizeof(*out)); + int rc = TF_INVALID; + mem_t mem = {data, len, 0}; + AVFormatContext* fmt = NULL; + AVIOContext* io = NULL; + AVCodecContext* dec = NULL; + AVPacket* pkt = NULL; + AVFrame* frame = NULL; + struct SwsContext* sws = NULL; + int* wanted = NULL; + unsigned char* iobuf = av_malloc(1 << 16); + if (!iobuf) return TF_MEMORY; + io = avio_alloc_context(iobuf, 1 << 16, 0, &mem, mem_read, NULL, mem_seek); + if (!io) { av_free(iobuf); return TF_MEMORY; } + fmt = avformat_alloc_context(); + if (!fmt) { rc = TF_MEMORY; goto done; } + fmt->pb = io; + if (avformat_open_input(&fmt, NULL, NULL, NULL) < 0) { fmt = NULL; goto done; } + if (avformat_find_stream_info(fmt, NULL) < 0) goto done; + int si = av_find_best_stream(fmt, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0); + if (si < 0) { rc = TF_NO_STREAM; goto done; } + // PyAV's streams.video[0] is the first video stream + for (unsigned i = 0; i < fmt->nb_streams; ++i) + if (fmt->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { si = (int)i; break; } + AVStream* st = fmt->streams[si]; + int width = st->codecpar->width, height = st->codecpar->height; + if (width <= 0 || height <= 0 || width > max_dim || height > max_dim) { rc = TF_DIMENSIONS; goto done; } + AVRational r = st->avg_frame_rate; + if (r.num <= 0 || r.den <= 0) r = av_guess_frame_rate(fmt, st, NULL); + double fps = (r.num > 0 && r.den > 0) ? av_q2d(r) : 24.0; + double seconds = 0.0; + if (st->duration > 0 && st->duration != AV_NOPTS_VALUE) seconds = (double)st->duration * av_q2d(st->time_base); + else if (fmt->duration > 0 && fmt->duration != AV_NOPTS_VALUE) seconds = (double)fmt->duration / 1e6; + int64_t total = st->nb_frames > 0 ? st->nb_frames : (int64_t)round_half_even(seconds * fps); + if (!(fps > 0 && fps <= 1000) || total <= 0) { rc = TF_NO_COUNT; goto done; } + if ((double)total / fps > max_seconds) { rc = TF_TOO_LONG; goto done; } + int64_t count = (int64_t)((double)total / fps * rate); + if (count < min_frames) count = min_frames; + if (count > max_frames) count = max_frames; + if (count > total) count = total; + wanted = malloc(sizeof(int) * (size_t)count); + if (!wanted) { rc = TF_MEMORY; goto done; } + // numpy.linspace(0, total - 1, count): step = (stop - start) / (count - 1), the last exactly stop + double stop = (double)(total - 1); + for (int64_t i = 0; i < count; ++i) { + double v = count > 1 ? (double)i * (stop / (double)(count - 1)) : 0.0; + if (count > 1 && i == count - 1) v = stop; + wanted[i] = (int)round_half_even(v); + } + int oh = 0, ow = 0; + if (size(ctx, (int)count, height, width, &oh, &ow) != 0 || oh <= 0 || ow <= 0) { rc = TF_SIZE; goto done; } + const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id); + if (!codec) goto done; + dec = avcodec_alloc_context3(codec); + if (!dec || avcodec_parameters_to_context(dec, st->codecpar) < 0) goto done; + dec->thread_count = 0; // PyAV's thread_type "AUTO" + dec->thread_type = FF_THREAD_FRAME | FF_THREAD_SLICE; + if (avcodec_open2(dec, codec, NULL) < 0) goto done; + out->rgb = malloc((size_t)count * (size_t)oh * (size_t)ow * 3); + out->indices = malloc(sizeof(int) * (size_t)count); + pkt = av_packet_alloc(); + frame = av_frame_alloc(); + if (!out->rgb || !out->indices || !pkt || !frame) { rc = TF_MEMORY; goto done; } + AVFrame* rgb = av_frame_alloc(); + if (!rgb) { rc = TF_MEMORY; goto done; } + sws = sws_alloc_context(); + if (!sws) { av_frame_free(&rgb); rc = TF_MEMORY; goto done; } + sws->threads = 0; + sws->flags = SWS_BICUBIC; + int n = 0, next = 0, got = 0, last = wanted[count - 1]; + int flushing = 0, finished = 0; + while (!finished) { + if (!flushing) { + int e = av_read_frame(fmt, pkt); + if (e < 0) { flushing = 1; avcodec_send_packet(dec, NULL); } + else { + if (pkt->stream_index == si) avcodec_send_packet(dec, pkt); + av_packet_unref(pkt); + } + } + for (;;) { + int e = avcodec_receive_frame(dec, frame); + if (e == AVERROR(EAGAIN)) { + if (flushing) finished = 1; // a flushed decoder only ends + break; + } + if (e < 0) { finished = 1; break; } + if (n > last) { finished = 1; av_frame_unref(frame); break; } + if (next < count && n == wanted[next]) { + // PyAV's _reformat: the frame's colour space and range kept; transfer and primaries unspecified + av_frame_unref(rgb); + av_frame_copy_props(rgb, frame); + set_colorspace(frame, (int)frame->colorspace); + rgb->colorspace = frame->colorspace; + rgb->color_range = frame->color_range; + set_colorspace(rgb, (int)frame->colorspace); + const enum AVColorTransferCharacteristic trc = frame->color_trc; + const enum AVColorPrimaries pri = frame->color_primaries; + frame->color_trc = rgb->color_trc = AVCOL_TRC_UNSPECIFIED; + frame->color_primaries = rgb->color_primaries = AVCOL_PRI_UNSPECIFIED; + rgb->format = AV_PIX_FMT_RGB24; + rgb->width = ow; + rgb->height = oh; + if (av_frame_get_buffer(rgb, 0) < 0) { rc = TF_MEMORY; av_frame_unref(frame); av_frame_free(&rgb); goto done; } + int e2 = sws_scale_frame(sws, rgb, frame); + frame->color_trc = trc; + frame->color_primaries = pri; + if (e2 < 0) { av_frame_unref(frame); av_frame_free(&rgb); goto done; } + uint8_t* to = out->rgb + (size_t)got * oh * ow * 3; + for (int y = 0; y < oh; ++y) memcpy(to + (size_t)y * ow * 3, rgb->data[0] + (size_t)y * rgb->linesize[0], (size_t)ow * 3); + out->indices[got++] = n; + while (next < count && wanted[next] == n) ++next; + } + ++n; + av_frame_unref(frame); + } + } + av_frame_free(&rgb); + if (got < 2) { rc = TF_FEW_FRAMES; goto done; } + out->frames = got; + out->height = oh; + out->width = ow; + out->fps = fps; + rc = TF_OK; +done: + if (rc != TF_OK) tf_video_free(out); + free(wanted); + sws_freeContext(sws); + av_frame_free(&frame); + av_packet_free(&pkt); + avcodec_free_context(&dec); + if (fmt) avformat_close_input(&fmt); + if (io) { av_freep(&io->buffer); avio_context_free(&io); } + return rc; +} diff --git a/zig/tests/bolt_vision_main.zig b/zig/tests/bolt_vision_main.zig new file mode 100644 index 0000000000..d64e3900be --- /dev/null +++ b/zig/tests/bolt_vision_main.zig @@ -0,0 +1,177 @@ +//! `tf-bolt-vision [reps]`: Bolt's image path against bolt_vision_reference.py's --dump. +const std = @import("std"); +const cuda = @import("cuda"); +const nemotron = @import("nemotron"); +const qi = @import("qwen_image"); + +const V = nemotron.vision; +const NI = qi.nemotron; + +pub fn main(init: std.process.Init) !u8 { + const a = init.arena.allocator(); + f32_arena = a; + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len < 3 or args.len > 4) { + std.debug.print("usage: tf-bolt-vision [reps]\n", .{}); + return 2; + } + const model = args[1]; + const dir = args[2]; + const reps: usize = if (args.len == 4) try std.fmt.parseInt(usize, args[3], 10) else 3; + const cwd = std.Io.Dir.cwd(); + + // host: decode is the shared decoder's; the reference's RGB is the input here + const rgb = try readNpy(a, io, dir, "image_rgb.npy"); + if (rgb.shape.len != 3 or rgb.shape[2] != 3 or !std.mem.eql(u8, rgb.descr, "|u1")) return error.BadFixture; + const im: NI.Image = .{ .w = @intCast(rgb.shape[1]), .h = @intCast(rgb.shape[0]), .rgb = @constCast(rgb.data) }; + const meta = try std.json.parseFromSliceLeaky(std.json.Value, a, try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "meta.json" }), a, .limited(1 << 20)), .{}); + const ghw = meta.object.get("grid_hw").?.array.items; + var t0 = now(io); + const g = try NI.grid(im.w, im.h); + const px = try NI.pixelValues(a, im, g[1] * NI.PATCH, g[0] * NI.PATCH); + const patches = try NI.patchify(a, px, g[0], g[1]); + const host_ms = ms(io, t0); + std.debug.print("image {d}x{d} -> grid {d}x{d} (reference {d}x{d}), {d} patches, {d} image tokens; host prep {d:.1} ms\n", .{ im.w, im.h, g[0], g[1], ghw[0].integer, ghw[1].integer, @as(usize, g[0]) * g[1], @as(usize, g[0]) * g[1] / 4, host_ms }); + if (g[0] != ghw[0].integer or g[1] != ghw[1].integer) return 1; + var fail = false; + { + const want = try readNpy(a, io, dir, "pixel_values.npy"); + const st = compare(f32Slice(want.data), px, 3 * @as(usize, g[1]) * NI.PATCH); + report("pixel_values", st); + if (st.max_abs > 1e-4) fail = true; + } + + var driver = try cuda.Driver.open(); + defer driver.close(); + var ctx = try cuda.Context.init(&driver, 0); + defer ctx.deinit(); + var stream = try cuda.Stream.init(ctx.d, true); + defer stream.deinit(); + const before = try ctx.memInfo(); + t0 = now(io); + const tower = try V.Tower.load(init.gpa, io, ctx.d, model); + defer tower.deinit(); + const after = try ctx.memInfo(); + std.debug.print("loaded in {d:.0} ms: arena {d:.3} GB (weights {d:.3} GB, scratch for {d} patches {d:.3} GB); device use +{d:.3} GB\n", .{ ms(io, t0), gb(tower.arena.len), gb(V.Tower.weightBytes()), V.max_patches, gb(V.Tower.scratchBytes(V.max_patches)), gb(before.free -| after.free) }); + + const n = @as(usize, g[0]) * g[1]; + const rows = V.prefix_tokens + n; + const tokens = n / 4; + var out = try cuda.DeviceBuffer.alloc(ctx.d, tokens * V.out_hidden * 2); + defer out.free(); + const xs = try a.alloc(f32, rows * V.hidden); + + // the embedded rows: no blocks run + tower.run_depth = 0; + var stages = try tower.encode(stream, patches, g[0], g[1], out.ptr); + try dl(ctx.d, stages.x, std.mem.sliceAsBytes(xs)); + report("tower_embed", compare(f32Slice((try readNpy(a, io, dir, "tower_embed.npy")).data), xs, V.hidden)); + tower.run_depth = V.depth; + + var times: std.ArrayList(f64) = .empty; + for (0..@max(reps, 1)) |_| { + t0 = now(io); + stages = try tower.encode(stream, patches, g[0], g[1], out.ptr); + try times.append(a, ms(io, t0)); + } + try dl(ctx.d, stages.x, std.mem.sliceAsBytes(xs)); + const tower_want = f32Slice((try readNpy(a, io, dir, "tower_out.npy")).data); + report("tower_out (all rows)", compare(tower_want, xs, V.hidden)); + report("tower_features (patch rows)", compare(tower_want[V.prefix_tokens * V.hidden ..], xs[V.prefix_tokens * V.hidden ..], V.hidden)); + const got16 = try a.alloc(u16, tokens * V.out_hidden); + try out.download(0, std.mem.sliceAsBytes(got16)); + const got = try a.alloc(f32, got16.len); + for (got, got16) |*o, h| o.* = @bitCast(@as(u32, h) << 16); + const proj = compare(f32Slice((try readNpy(a, io, dir, "projector_out.npy")).data), got, V.out_hidden); + report("projector_out", proj); + if (proj.cosine < 0.999) fail = true; + std.debug.print("encode ms ({d} patches, {d} tokens; first call makes the GEMM plans):", .{ n, tokens }); + for (times.items) |t| std.debug.print(" {d:.1}", .{t}); + std.debug.print("\n{s}\n", .{if (fail) "FAIL" else "PASS"}); + return if (fail) 1 else 0; +} + +const Stats = struct { max_abs: f64, max_ref: f64, rel: f64, cosine: f64, row_min: f64, rows: usize }; + +/// `got` against `want`, both row-major with rows of `width`. +fn compare(want: []const f32, got: []const f32, width: usize) Stats { + std.debug.assert(want.len == got.len); + var s: Stats = .{ .max_abs = 0, .max_ref = 0, .rel = 0, .cosine = 0, .row_min = 1, .rows = got.len / width }; + var dot: f64 = 0; + var nw: f64 = 0; + var ng: f64 = 0; + var dd: f64 = 0; + for (0..s.rows) |r| { + var rd: f64 = 0; + var rw: f64 = 0; + var rg: f64 = 0; + for (want[r * width ..][0..width], got[r * width ..][0..width]) |w0, g0| { + const w: f64 = w0; + const gg: f64 = g0; + s.max_abs = @max(s.max_abs, @abs(w - gg)); + s.max_ref = @max(s.max_ref, @abs(w)); + rd += w * gg; + rw += w * w; + rg += gg * gg; + dd += (w - gg) * (w - gg); + } + dot += rd; + nw += rw; + ng += rg; + if (rw > 0 and rg > 0) s.row_min = @min(s.row_min, rd / @sqrt(rw * rg)); + } + s.cosine = dot / @sqrt(nw * ng); + s.rel = @sqrt(dd / nw); + return s; +} + +fn report(name: []const u8, s: Stats) void { + std.debug.print("{s}: {d} rows, max|d| {e:.3} (max|ref| {d:.3}), rel {e:.3}, cosine {d:.6}, worst row {d:.6}\n", .{ name, s.rows, s.max_abs, s.max_ref, s.rel, s.cosine, s.row_min }); +} + +const Npy = struct { descr: []const u8, shape: []usize, data: []const u8 }; + +/// A little-endian C-order .npy (version 1 or 2 header). +fn readNpy(a: std.mem.Allocator, io: std.Io, dir: []const u8, name: []const u8) !Npy { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, try std.fs.path.join(a, &.{ dir, name }), a, .limited(1 << 31)); + if (bytes.len < 10 or !std.mem.eql(u8, bytes[0..6], "\x93NUMPY")) return error.BadNpy; + const v2 = bytes[6] >= 2; + const hl: usize = if (v2) std.mem.readInt(u32, bytes[8..12], .little) else std.mem.readInt(u16, bytes[8..10], .little); + const start: usize = if (v2) 12 else 10; + const header = bytes[start..][0..hl]; + if (std.mem.indexOf(u8, header, "'fortran_order': True") != null) return error.BadNpy; + const d0 = (std.mem.indexOf(u8, header, "'descr': '") orelse return error.BadNpy) + 10; + const d1 = std.mem.indexOfScalarPos(u8, header, d0, '\'') orelse return error.BadNpy; + const s0 = (std.mem.indexOf(u8, header, "'shape': (") orelse return error.BadNpy) + 10; + const s1 = std.mem.indexOfScalarPos(u8, header, s0, ')') orelse return error.BadNpy; + var shape: std.ArrayList(usize) = .empty; + var it = std.mem.tokenizeAny(u8, header[s0..s1], ", "); + while (it.next()) |tok| try shape.append(a, try std.fmt.parseInt(usize, tok, 10)); + return .{ .descr = header[d0..d1], .shape = shape.items, .data = bytes[start + hl ..] }; +} + +var f32_arena: std.mem.Allocator = undefined; + +/// The npy payload as f32 (copied: the file's bytes need not be 4-aligned in memory). +fn f32Slice(b: []const u8) []const f32 { + const out = f32_arena.alloc(f32, b.len / 4) catch @panic("out of memory"); + @memcpy(std.mem.sliceAsBytes(out), b[0 .. out.len * 4]); + return out; +} + +fn dl(d: *const cuda.Driver, src: u64, out: []u8) !void { + try d.check(d.api.cuMemcpyDtoH_v2(out.ptr, src, out.len), "cuMemcpyDtoH"); +} + +fn now(io: std.Io) i96 { + return std.Io.Clock.awake.now(io).toNanoseconds(); +} + +fn ms(io: std.Io, t0: i96) f64 { + return @as(f64, @floatFromInt(now(io) - t0)) / 1e6; +} + +fn gb(b: usize) f64 { + return @as(f64, @floatFromInt(b)) / 1e9; +} diff --git a/zig/tests/cuda/ggml_tests.zig b/zig/tests/cuda/ggml_tests.zig new file mode 100644 index 0000000000..849349b707 --- /dev/null +++ b/zig/tests/cuda/ggml_tests.zig @@ -0,0 +1,336 @@ +//! Bolt's ggml formats on the GPU (ggml.cu): real tensors' repack, grouped IQ2_XXS experts and dense rows against fp64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const Fixture = @import("fixture.zig").Fixture; +const Gpu = check.Gpu; + +const ggml = cuda.ggml; +const grouped = cuda.grouped; +const ex_ = cuda.experts; + +fn arr(fx: Fixture, comptime fmt: []const u8, args: anytype) ![]u8 { + var name: [48]u8 = undefined; + return fx.bytes(try std.fmt.bufPrint(&name, fmt, args)); +} + +fn bf(v: u16) f64 { + return @floatCast(@as(f32, @bitCast(@as(u32, v) << 16))); +} + +fn relu2(v: f64) f64 { + const u = @max(v, 0); + return u * u; +} + +/// Max |got - want| and max |want| over outputs fp32 or bf16; fails past `tol` of the largest reference. +fn compare(gpu: Gpu, what: []const u8, b: cuda.DeviceBuffer, ref: []const u8, out: ggml.Out, tol: f64) !void { + const got = try check.download(gpu, b); + defer gpu.gpa.free(got); + const want = std.mem.bytesAsSlice(f64, @as([]align(1) const u8, ref)); + var err: f64 = 0; + var top: f64 = 0; + for (want, 0..) |w0, i| { + const w = if (out == .relu2) relu2(w0) else w0; + const g: f64 = if (out == .f32) @floatCast(std.mem.bytesToValue(f32, got[4 * i ..][0..4])) else bf(std.mem.bytesToValue(u16, got[2 * i ..][0..2])); + if (!(g == g)) return check.expect(false, "{s}: NaN at {d}", .{ what, i }); + err = @max(err, @abs(g - w)); + top = @max(top, @abs(w)); + } + try check.expect(err <= tol * top, "{s}: max abs error {e:.3} is {e:.3} of max |ref| {e:.3}", .{ what, err, err / top, top }); + check.pass("{s}: max abs error {e:.3}, {e:.3} of max |ref| {e:.3}", .{ what, err, err / top, top }); +} + +/// fp32 sums; bf16 rounds once (2^-8 relative); relu^2 rounds its input to bf16 too, then squares it. +fn tolerance(out: ggml.Out) f64 { + return switch (out) { + .f32 => 1e-5, + .bf16 => 1.0 / 256.0, + .relu2 => 3.0 / 256.0, + }; +} + +/// Raw blocks packed on the device equal the host repack. +fn packed_(gpu: Gpu, k: *const ggml.Kernels, s: cuda.Stream, kind: ggml.Kind, raw: []const u8, n: usize, kk: usize, count: usize) !cuda.DeviceBuffer { + var draw = try cuda.DeviceBuffer.fromHost(gpu.d, raw); + defer draw.free(); + var out = try cuda.DeviceBuffer.alloc(gpu.d, raw.len); + errdefer out.free(); + try k.pack(s, kind, draw.ptr, out.ptr, n, kk, count); + try s.synchronize(); + const got = try check.download(gpu, out); + defer gpu.gpa.free(got); + const want = try gpu.gpa.alloc(u8, raw.len); + defer gpu.gpa.free(want); + const one = ggml.bytes(kind, n, kk); + for (0..count) |m| ggml.packHost(kind, want[m * one ..][0..one], raw[m * one ..][0..one], n, kk); + try check.sameBytes("ggml pack", got, want); + return out; +} + +pub fn run(gpu: Gpu, dir: []const u8) !void { + var fx = try Fixture.open(gpu.gpa, gpu.io, dir); + defer fx.deinit(); + var k = try ggml.Kernels.load(gpu.d); + defer k.unload(); + var stream = try cuda.Stream.init(gpu.d, false); // blocking: launches wait for fromHost's legacy-stream copies + defer stream.deinit(); + try experts(gpu, fx, &k, stream); + try dense(gpu, fx, &k, stream); +} + +fn experts(gpu: Gpu, fx: Fixture, k: *const ggml.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + const e: usize = @intCast(try fx.int("experts")); + const d: usize = @intCast(try fx.int("dims")); + const ni: usize = @intCast(try fx.int("width")); + const slots: usize = @intCast(try fx.int("slots")); + const up_raw = try fx.bytes("up_raw"); + defer a.free(up_raw); + const down_raw = try fx.bytes("down_raw"); + defer a.free(down_raw); + var up = try packed_(gpu, k, s, .iq2xxs, up_raw, ni, d, e); + defer up.free(); + var down = try packed_(gpu, k, s, .iq2xxs, down_raw, d, ni, e); + defer down.free(); + check.pass("ggml pack: {d} real IQ2_XXS experts' up [{d}, {d}] and down on the device equal the host repack", .{ e, ni, d }); + const layer: ex_.Layer = .{ .format = .iq2xxs, .up = up.ptr, .down = down.ptr, .width = ni, .dims = d, .experts = e, .relu2 = true }; + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + const g: ex_.Grouped = .{ .iq2xxs = k }; + inline for (.{ .{ "decode", false }, .{ "prompt", true } }) |set| { + var it = std.mem.tokenizeScalar(u8, try fx.string(set[0]), ','); + while (it.next()) |tok| { + const rows = try std.fmt.parseInt(usize, tok, 10); + try expertCase(gpu, fx, router, g, k, s, layer, rows, slots, set[1]); + if (set[1]) try expertCase(gpu, fx, router, g, k, s, layer, rows, slots, false); // decode form at prompt width + } + } +} + +fn expertCase(gpu: Gpu, fx: Fixture, router: grouped.Router, g: ex_.Grouped, k: *const ggml.Kernels, s: cuda.Stream, l: ex_.Layer, rows: usize, slots: usize, prompt: bool) !void { + const a = gpu.gpa; + const pairs = rows * slots; + const tile: usize = if (prompt) ex_.promptTile(.iq2xxs) else 16; + const picks = try arr(fx, "picks{d}", .{rows}); + defer a.free(picks); + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, picks); + defer dpicks.free(); + const big = pairs > grouped.small; + const bound = grouped.maxItems(pairs, l.experts, tile); + var bufs: [5]cuda.DeviceBuffer = undefined; + const sizes = [5]usize{ pairs * 4, bound * 12, 8, (if (big) pairs else 1) * 4, (if (big) (pairs + 1023) / 1024 * l.experts else 1) * 4 }; + var made: usize = 0; + defer for (bufs[0..made]) |*b| b.free(); + for (sizes, &bufs) |n, *b| { + b.* = try cuda.DeviceBuffer.alloc(gpu.d, n); + made += 1; + } + const plan: grouped.Plan = .{ .members = bufs[0].ptr, .items = bufs[1].ptr, .counts = bufs[2].ptr, .rank = bufs[3].ptr, .hist = bufs[4].ptr }; + try router.route(s, dpicks.ptr, pairs, l.experts, tile, plan); + + const x = try arr(fx, "x{d}", .{rows}); + defer a.free(x); + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, x); + defer dx.free(); + const h = try arr(fx, "h{d}", .{rows}); + defer a.free(h); + var dh = try cuda.DeviceBuffer.fromHost(gpu.d, h); + defer dh.free(); + const ref_up = try arr(fx, "up{d}", .{rows}); + defer a.free(ref_up); + const ref_down = try arr(fx, "down{d}", .{rows}); + defer a.free(ref_down); + var y = try cuda.DeviceBuffer.alloc(gpu.d, pairs * @max(l.width, l.dims) * 4); + defer y.free(); + const form = if (prompt) "prompt" else "decode"; + var name: [96]u8 = undefined; + + // up: relu^2 through experts.Grouped, then fp32 sums by the same kernels with no activation + try y.fill8(0xff, s.handle); // a call that writes nothing fails rather than pass on old bytes + const call = if (prompt) &ex_.Grouped.prompt else &ex_.Grouped.decode; + try call(g, s, .up, .{ .x = dx.ptr, .stride = l.dims, .slots = slots }, l, plan, y.ptr, bound, -1); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "iq2 up relu^2 {s} {d} rows x {d}", .{ form, rows, slots }), y, ref_up, .relu2, tolerance(.relu2)); + try y.fill8(0xff, s.handle); + try k.experts(s, prompt, .iq2xxs, .f32, dx.ptr, l.dims, slots, l.up, l.width, l.dims, plan, y.ptr, bound, -1); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "iq2 up fp32 {s} {d} rows x {d}", .{ form, rows, slots }), y, ref_up, .f32, tolerance(.f32)); + // down from pair rows, fp32 and bf16 + inline for (.{ .{ ex_.Epi.down_f32, ggml.Out.f32 }, .{ ex_.Epi.down_bf16, ggml.Out.bf16 } }) |c| { + try y.fill8(0xff, s.handle); + try call(g, s, c[0], .{ .x = dh.ptr, .stride = l.width }, l, plan, y.ptr, bound, -1); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "iq2 down {t} {s} {d} rows x {d}", .{ c[1], form, rows, slots }), y, ref_down, c[1], tolerance(c[1])); + } +} + +fn dense(gpu: Gpu, fx: Fixture, k: *const ggml.Kernels, s: cuda.Stream) !void { + const a = gpu.gpa; + var cases = std.mem.tokenizeScalar(u8, try fx.string("dense"), ','); + while (cases.next()) |case| { + var name: [96]u8 = undefined; + const n: usize = @intCast(try fx.int(try std.fmt.bufPrint(&name, "{s}_n", .{case}))); + const kk: usize = @intCast(try fx.int(try std.fmt.bufPrint(&name, "{s}_k", .{case}))); + const kind: ggml.Kind = if (std.mem.eql(u8, try fx.string(try std.fmt.bufPrint(&name, "{s}_kind", .{case})), "Q8_0")) .q8_0 else .iq2xxs; + const raw = try arr(fx, "{s}_raw", .{case}); + defer a.free(raw); + var w_buf = try packed_(gpu, k, s, kind, raw, n, kk, 1); + defer w_buf.free(); + const w = ggml.Weight{ .kind = kind, .w = w_buf.ptr, .n = n, .k = kk }; + var rows_it = std.mem.tokenizeScalar(u8, try fx.string("dense_rows"), ','); + while (rows_it.next()) |tok| { + const rows = try std.fmt.parseInt(usize, tok, 10); + const x = try arr(fx, "{s}_x{d}", .{ case, rows }); + defer a.free(x); + var dx = try cuda.DeviceBuffer.fromHost(gpu.d, x); + defer dx.free(); + const ref = try arr(fx, "{s}_y{d}", .{ case, rows }); + defer a.free(ref); + var y = try cuda.DeviceBuffer.alloc(gpu.d, rows * n * 4); + defer y.free(); + const forms: []const bool = if (rows <= ggml.decode_rows) &.{ false, true } else &.{true}; + for (forms) |prompt| for ([_]ggml.Out{ .f32, .bf16 }) |out| { + try y.fill8(0xff, s.handle); + if (prompt) try k.promptRows(s, w, dx.ptr, 0, rows, y.ptr, out) else try k.decode(s, w, dx.ptr, 0, rows, y.ptr, out); + try s.synchronize(); + try compare(gpu, try std.fmt.bufPrint(&name, "{t} {s} [{d}, {d}] {s} {d} rows {t}", .{ kind, case, n, kk, if (prompt) "prompt" else "decode", rows, out }), y, ref, out, tolerance(out)); + }; + } + } +} + +/// Decode GB/s: 9 IQ2_XXS experts' up and down at width 2560 (fresh experts each pass, past L2), and lm_head rows. +pub fn bench(gpu: Gpu) !void { + var k = try ggml.Kernels.load(gpu.d); + defer k.unload(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try grouped.Router.resolve(mod); + var t0 = try cuda.Event.init(gpu.d, true); + defer t0.deinit(); + var t1 = try cuda.Event.init(gpu.d, true); + defer t1.deinit(); + + // experts: 16 sets of 9 (144 experts, 194 MB), a pass on each set in turn + const width = 2560; + const dims = 1024; + const sets = 16; + const e = 9 * sets; + const one = ggml.bytes(.iq2xxs, width, dims); + var up = try cuda.DeviceBuffer.alloc(gpu.d, e * one); + defer up.free(); + var down = try cuda.DeviceBuffer.alloc(gpu.d, e * one); + defer down.free(); + try up.fill8(0x11, s.handle); + try down.fill8(0x11, s.handle); + var x = try cuda.DeviceBuffer.alloc(gpu.d, 512 * 4096 * 2); + defer x.free(); + try x.fill8(0x3c, s.handle); // bf16 ~0.01 + var h = try cuda.DeviceBuffer.alloc(gpu.d, 9 * width * 2); + defer h.free(); + try h.fill8(0x3c, s.handle); + var y = try cuda.DeviceBuffer.alloc(gpu.d, 512 * 131072 * 2); + defer y.free(); + var plans = try cuda.DeviceBuffer.alloc(gpu.d, sets * 256); + defer plans.free(); + var picks: [sets * 9]i32 = undefined; + for (&picks, 0..) |*p, i| p.* = @intCast(i); + var dpicks = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&picks)); + defer dpicks.free(); + var ps: [sets]grouped.Plan = undefined; + for (&ps, 0..) |*p, i| { + const base = plans.ptr + i * 256; + p.* = .{ .members = base, .items = base + 64, .counts = base + 192, .rank = 0, .hist = 0 }; + try router.route(s, dpicks.ptr + i * 9 * 4, 9, e, 16, p.*); + } + const l: ex_.Layer = .{ .format = .iq2xxs, .up = up.ptr, .down = down.ptr, .width = width, .dims = dims, .experts = e, .relu2 = true }; + const g: ex_.Grouped = .{ .iq2xxs = &k }; + const reps = 20; + const bound = grouped.maxItems(9, e, 16); + const parts = [_][]const u8{ "up", "down", "up+down" }; + for (parts, 0..) |what, pi| { + for (0..2) |warm| { + try t0.record(s); + for (0..reps) |_| for (ps) |p| { + if (pi != 1) try g.decode(s, .up, .{ .x = x.ptr, .stride = dims, .slots = 9 }, l, p, h.ptr, bound, -1); + if (pi != 0) try g.decode(s, .down_f32, .{ .x = h.ptr, .stride = width }, l, p, y.ptr, bound, -1); + }; + try t1.record(s); + try t1.synchronize(); + if (warm == 0) continue; + const us = try cuda.Event.elapsedMs(t0, t1) * 1e3 / (reps * sets); + const moved: f64 = @floatFromInt(9 * one * @as(usize, if (pi == 2) 2 else 1)); + std.debug.print("RESULT iq2xxs experts decode, 1 token x 9 experts {s} at width {d}: {d:.2} us, {d:.0} GB/s\n", .{ what, width, us, moved / us / 1e3 }); + } + } + + // prompt rows: 2048 tokens x 9 slots over the 144 experts, up then down, as the prompt plan groups them + { + const tokens = 2048; + const pairs = tokens * 9; + const many = try gpu.gpa.alloc(i32, pairs); + defer gpu.gpa.free(many); + var prng = std.Random.DefaultPrng.init(7); + for (0..tokens) |tk| { + const first = prng.random().uintLessThan(usize, e - 9); + for (0..9) |j| many[tk * 9 + j] = @intCast(first + j); + } + var dmany = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(many)); + defer dmany.free(); + const pb = grouped.maxItems(pairs, e, 64); + const hist = (pairs + 1023) / 1024 * e * 4; + var scratch = try cuda.DeviceBuffer.alloc(gpu.d, pairs * 8 + pb * 12 + 16 + hist); + defer scratch.free(); + const base = scratch.ptr; + const pp: grouped.Plan = .{ .members = base, .items = base + pairs * 4, .counts = base + pairs * 4 + pb * 12, .rank = base + pairs * 4 + pb * 12 + 16, .hist = base + pairs * 8 + pb * 12 + 16 }; + try router.route(s, dmany.ptr, pairs, e, 64, pp); + var xp = try cuda.DeviceBuffer.alloc(gpu.d, tokens * dims * 2); + defer xp.free(); + try xp.fill8(0x3c, s.handle); + var hp = try cuda.DeviceBuffer.alloc(gpu.d, pairs * width * 2); + defer hp.free(); + var yp = try cuda.DeviceBuffer.alloc(gpu.d, pairs * dims * 2); + defer yp.free(); + var t2 = try cuda.Event.init(gpu.d, true); + defer t2.deinit(); + for (0..2) |warm| { + try t0.record(s); + try g.prompt(s, .up, .{ .x = xp.ptr, .stride = dims, .slots = 9 }, l, pp, hp.ptr, pb, -1); + try t1.record(s); + try g.prompt(s, .down_bf16, .{ .x = hp.ptr, .stride = width }, l, pp, yp.ptr, pb, -1); + try t2.record(s); + try t2.synchronize(); + if (warm == 0) continue; + const up_us = try cuda.Event.elapsedMs(t0, t1) * 1e3; + const down_us = try cuda.Event.elapsedMs(t1, t2) * 1e3; + const flops: f64 = 2.0 * pairs * width * dims; + std.debug.print("RESULT iq2xxs experts prompt, {d} tokens x 9 over {d} experts at width {d}: up {d:.0} us ({d:.0} TFLOP/s), down {d:.0} us ({d:.0} TFLOP/s)\n", .{ tokens, e, width, up_us, flops / up_us / 1e6, down_us, flops / down_us / 1e6 }); + } + } + + // lm_head: Q8_0 [131072, 4096] + const n = 131072; + const kk = 4096; + const lb = ggml.bytes(.q8_0, n, kk); + var lm = try cuda.DeviceBuffer.alloc(gpu.d, lb); + defer lm.free(); + try lm.fill8(0x11, s.handle); + const w = ggml.q8Weight(lm.ptr, n, kk); + for ([_]usize{ 1, 4, 16, 64, 512 }) |rows| for ([_]ggml.Out{ .f32, .bf16 }) |out| { + if (rows > 16 and out == .f32) continue; + const r2: usize = if (rows > 16) 3 else 10; + for (0..2) |warm| { + try t0.record(s); + for (0..r2) |_| try k.matmul(s, w, x.ptr, 0, rows, y.ptr, out); + try t1.record(s); + try t1.synchronize(); + if (warm == 0) continue; + const us = try cuda.Event.elapsedMs(t0, t1) * 1e3 / @as(f32, @floatFromInt(r2)); + std.debug.print("RESULT q8_0 lm_head [{d}, {d}] {d} rows {t}: {d:.1} us, {d:.0} GB/s\n", .{ n, kk, rows, out, us, @as(f64, @floatFromInt(lb)) / us / 1e3 }); + } + }; +} diff --git a/zig/tests/cuda/glue_tests.zig b/zig/tests/cuda/glue_tests.zig index 7af2fd42b9..331398e564 100644 --- a/zig/tests/cuda/glue_tests.zig +++ b/zig/tests/cuda/glue_tests.zig @@ -177,20 +177,25 @@ fn norms(r: *Rig) !void { } fn route(r: *Rig) !void { - const sk = glue.routerShape(D).sk; + try routeShape(r, D, E, top_k, NS); + try routeShape(r, 4096, 496, 9, 9); // Bolt: 4096-wide rows (eight K slices), 496 experts, top 9, no shared slots +} + +fn routeShape(r: *Rig, d: usize, e: usize, k: usize, ns: usize) !void { + const sk = glue.routerShape(d).sk; for ([_]usize{ 1, 16, 300 }) |rows| { - const x = try r.bfs(rows * D, 1.0); - const w = try r.bfs(E * D, 0.05); - const bias = try r.f32s(E, -0.05, 0.05); - const part = try r.a.alloc(f32, sk * rows * E); - const idx = try r.a.alloc(i32, rows * NS); - const wt = try r.a.alloc(f32, rows * NS); - ref.router(x, w, part, rows, D, E, sk); - ref.topk(part, bias, idx, wt, rows, 2.5, E, sk, top_k, NS, true); - const partd = try r.zeros(f32, sk * rows * E); - const idxd = try r.zeros(i32, rows * NS); - const wtd = try r.zeros(f32, rows * NS); - try r.g().route(try r.dev(u16, x), try r.dev(u16, w), try r.dev(f32, bias), partd, idxd, wtd, rows, D, E, top_k, 2.5, true); + const x = try r.bfs(rows * d, 1.0); + const w = try r.bfs(e * d, 0.05); + const bias = try r.f32s(e, -0.05, 0.05); + const part = try r.a.alloc(f32, sk * rows * e); + const idx = try r.a.alloc(i32, rows * ns); + const wt = try r.a.alloc(f32, rows * ns); + ref.router(x, w, part, rows, d, e, sk); + ref.topk(part, bias, idx, wt, rows, 2.5, e, sk, k, ns, true); + const partd = try r.zeros(f32, sk * rows * e); + const idxd = try r.zeros(i32, rows * ns); + const wtd = try r.zeros(f32, rows * ns); + try r.g().route(try r.dev(u16, x), try r.dev(u16, w), try r.dev(f32, bias), partd, idxd, wtd, rows, d, e, k, ns, 2.5, true); try r.same("router partials", f32, partd, part); try r.same("top-k ids", i32, idxd, idx); try r.same("top-k weights", f32, wtd, wt); diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index aec6b9ec1b..f8e95d6012 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -29,6 +29,7 @@ const norms_tests = @import("norms_tests.zig"); const norm_rows_tests = @import("norm_rows_tests.zig"); const pocket_tests = @import("pocket_tests.zig"); const recapture_tests = @import("recapture_tests.zig"); +const ggml_tests = @import("ggml_tests.zig"); const usage = \\usage: tf-cuda-test @@ -64,6 +65,8 @@ const usage = \\ quant-ref FP8 and NVFP4 projections and experts against float64 references (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes + \\ ggml Bolt's IQ2_XXS experts and Q8_0 / IQ2_XXS dense rows on real tensors against fp64 (tools/zig/bolt_ggml_fixture.py) + \\ ggml-bench IQ2_XXS expert and Q8_0 lm_head decode GB/s (synthetic blocks) \\ modelopt ModelOpt repack, FP8 and BF16 lane rows and relu^2 NVFP4 experts against fp64 \\ exlr8 EXLR8 projections: every codebook state, FP64, row windows, the Metal frozen digests \\ exlr8-bench DIR [MANIFEST] [LAYER] [SETS] a real EXLR8 layer: FP64, row windows, then GB/s by stage @@ -124,6 +127,8 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); + if (std.mem.eql(u8, cmd, "ggml")) return ggml_tests.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ggml-bench")) return ggml_tests.bench(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); if (std.mem.eql(u8, cmd, "exlr8-bench")) return exlr8_tests.bench(gpu, try arg(rest, 0), if (rest.len > 1) rest[1] else "home", if (rest.len > 2) try std.fmt.parseInt(u16, rest[2], 10) else 40, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 4);