llama: a smaller Llama drafting for the served one (--drafter), replies unchanged (builds on #1415)

docs/kernel-contract.md

@@ -141,6 +141,10 @@ Metal also checks Nemotron's routed experts on the engine's route tables, one ro
 against float64 at tf_load16's bf16 x sums (up's relu2 as an interval over its two bf16 roundings), and the rows2, rows4
 and six geometry pipelines of each projection bit-equal to them (two controls: member rows read from one row, and
 w3's x sums without their bf16 steps).
+The Llama family's own kernels (`metal_llama*.zig`): its glue (embedding, the norms at HF's rounding points, the
+residual add, RoPE with bf16 cos and sin, the SiLU gate, the decode attention's partials and their merge, the argmax
+draw by its stated rule), its row projections (1 to 32 rows) and, on the M5's tensor units, its four GEMM forms and the
+prompt attention at both served geometries; gate 5 reads their pipelines on this GPU.
 CUDA: the RMS norm rows (16-byte form and the fp32-out tree), the bf16 residual's norm tails (MLX's rounded rows and
 the once-rounded ones, each checked on its residual and its output, with four controls each dropping a rounding point), Nemotron's router logits (and their 128-row tiles,
 bit-equal to them) and top-k, the latent top-k's decode form (bit-equal to `tf_super_topk`), the latent combine and

docs/recipes/README.md

@@ -24,6 +24,7 @@ For agent-assisted contributions, read the [TensorFold contributor skill](SKILL.
 | Kolibri-1 | [Native CUDA, block FP8, research checkpoint](kolibri-1.md) |
 | LFM2.5-24B-A2B | [Native Metal, MLX 4-bit, shared rounds](lfm2.5-24b-a2b.md) |
 | Spark-X2.6-4B | [Native CUDA](spark-x2.6-4b-cuda.md), [native Metal, the bf16 release, contexts to 1M](spark-x2.6-4b-metal.md) |
+| Nanbeige4.3-33B | [Native Metal, the bf16 release, the Llama family](nanbeige-4.3-33b-metal.md) |
 
 ## Capability floor
 

tools/zig/llama_f64_reference.py

@@ -0,0 +1,97 @@
+"""A Llama checkpoint's teacher-forced logits in float64 on the CPU, rounded once to bf16: the exact reference."""
+# usage: llama_f64_reference.py MODEL_DIR tokens.json out.bin [--bf16-points]: raw bf16 [positions, vocab]
+import json
+import sys
+import time
+
+import numpy as np
+import torch
+from safetensors import safe_open
+
+
+def bf16_once(x: np.ndarray) -> np.ndarray:
+    """float64 to bf16 bits, rounded once to nearest (ties to even) at bf16's spacing."""
+    spacing = np.ldexp(1.0, np.frexp(np.abs(x))[1] - 1 - 7)
+    y = (np.round(x / spacing) * spacing).astype(np.float32)  # exact: y has at most 8 significant bits
+    return (y.view(np.uint32) >> 16).astype(np.uint16)
+
+
+def main(model_dir: str, tokens: str, out: str, points: bool = False) -> None:
+    """`points`: round to bf16 at every one of the reference's rounding points, each operation exact before it."""
+    torch.set_num_threads(24)
+    cfg = json.load(open(f"{model_dir}/config.json"))
+    index = json.load(open(f"{model_dir}/model.safetensors.index.json"))["weight_map"]
+    if cfg.get("rope_scaling") or cfg.get("attention_bias") or cfg.get("mlp_bias") or cfg.get("tie_word_embeddings"):
+        raise SystemExit("only the plain Llama layout: no scaled RoPE, biases or tied head")
+    layers, heads, kv = cfg["num_hidden_layers"], cfg["num_attention_heads"], cfg["num_key_value_heads"]
+    dim = cfg.get("head_dim") or cfg["hidden_size"] // heads
+    eps, theta = cfg["rms_norm_eps"], float(cfg["rope_theta"])
+    f64 = torch.float64
+    quant = cfg.get("quantization")
+    if quant and (quant.get("mode", "affine") != "affine" or quant["bits"] not in (4, 8)):
+        raise SystemExit("only MLX affine checkpoints at 4 or 8 bits")
+
+    def tensor(name):
+        with safe_open(f"{model_dir}/{index[name]}", framework="pt") as f:
+            return f.get_tensor(name)
+
+    def w(name, mlx=False):
+        if not name.endswith(".weight") or name[: -len("weight")] + "scales" not in index:
+            return tensor(name).to(f64)
+        # mlx_lm's layout: u32 words of `bits`-bit codes, lowest bits first; a scale and a bias a group of codes
+        bits, group = quant["bits"], quant["group_size"]
+        words = tensor(name).to(torch.int64) & 0xFFFFFFFF
+        shifts = torch.arange(0, 32, bits, dtype=torch.int64)
+        q = ((words.unsqueeze(-1) >> shifts) & ((1 << bits) - 1)).reshape(words.shape[0], -1).to(f64)
+        scales = tensor(name[: -len("weight")] + "scales").to(f64).repeat_interleave(group, -1)
+        biases = tensor(name[: -len("weight")] + "biases").to(f64).repeat_interleave(group, -1)
+        if mlx:  # mlx's dequantize, as an affine embedding's rows are read: bf16(bf16(scale * q) + bias)
+            b16 = lambda t: t.to(torch.bfloat16).to(f64)
+            return b16(b16(scales * q) + biases)
+        return scales * q + biases
+
+    def r(x):  # a rounding point: bf16 when `points`, else exact
+        return x.to(torch.bfloat16).to(f64) if points else x
+
+    def norm(x, g):
+        return r(g * r(x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)))
+
+    def rotate(x):
+        return torch.cat((-x[..., dim // 2 :], x[..., : dim // 2]), -1)
+
+    ids = json.load(open(tokens))
+    n = len(ids)
+    t0 = time.time()
+    with torch.no_grad():
+        inv = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=f64) / dim))
+        angle = torch.outer(torch.arange(n, dtype=f64), inv)
+        cos, sin = torch.cat((angle, angle), -1).cos(), torch.cat((angle, angle), -1).sin()
+        if points:  # HF's fp32 angle, its cos and sin rounded to bf16
+            a32 = torch.outer(torch.arange(n, dtype=torch.float32), 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)))
+            cos, sin = r(torch.cat((a32, a32), -1).cos().to(f64)), r(torch.cat((a32, a32), -1).sin().to(f64))
+        x = w("model.embed_tokens.weight", mlx=points)[torch.tensor(ids)]
+        mask = torch.full((n, n), float("-inf"), dtype=f64).triu(1)
+        for i in range(layers):
+            p = f"model.layers.{i}."
+            h = norm(x, w(p + "input_layernorm.weight"))
+            q = r(h @ w(p + "self_attn.q_proj.weight").T).view(n, heads, dim).transpose(0, 1)
+            k = r(h @ w(p + "self_attn.k_proj.weight").T).view(n, kv, dim).transpose(0, 1)
+            v = r(h @ w(p + "self_attn.v_proj.weight").T).view(n, kv, dim).transpose(0, 1)
+            q, k = r(r(q * cos) + r(rotate(q) * sin)), r(r(k * cos) + r(rotate(k) * sin))
+            k, v = k.repeat_interleave(heads // kv, 0), v.repeat_interleave(heads // kv, 0)
+            a = r(torch.softmax(q @ k.transpose(1, 2) / dim**0.5 + mask, -1) @ v)
+            x = r(x + r(a.transpose(0, 1).reshape(n, heads * dim) @ w(p + "self_attn.o_proj.weight").T))
+            h = norm(x, w(p + "post_attention_layernorm.weight"))
+            act = r(r(torch.nn.functional.silu(r(h @ w(p + "mlp.gate_proj.weight").T))) * r(h @ w(p + "mlp.up_proj.weight").T))
+            x = r(x + r(act @ w(p + "mlp.down_proj.weight").T))
+            print(f"layer {i + 1} of {layers}, {time.time() - t0:.0f} s", flush=True)
+        h = norm(x, w("model.norm.weight"))
+        head = w("lm_head.weight")
+        with open(out, "wb") as f:
+            for s in range(0, n, 64):
+                bf16_once((h[s : s + 64] @ head.T).numpy()).tofile(f)
+    print(f"{n} positions' logits in float64, rounded once to bf16, in {time.time() - t0:.0f} s")
+
+
+if __name__ == "__main__":
+    main(*sys.argv[1:4], points="--bf16-points" in sys.argv[4:])

tools/zig/llama_teacher_judge.py

@@ -0,0 +1,84 @@
+"""Judge a Llama checkpoint's teacher-forced logits against a reference's at every position."""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+from pathlib import Path
+
+from qwen27_margins import verdict
+
+
+def top2(row) -> dict:
+    """A row's best two ids and scores, ties to the lower id (as the engine's argmax); non-finite rows are refused."""
+    import numpy as np
+
+    if not np.isfinite(row).all():
+        raise ValueError(f"non-finite logits at id {int(np.flatnonzero(~np.isfinite(row))[0])}")
+    first = int(np.argmax(row))
+    rest = row.copy()
+    rest[first] = -np.inf
+    second = int(np.argmax(rest))
+    return {"ids": [first, second], "logits": [float(row[first]), float(row[second])]}
+
+
+def judge(native: list[dict], reference: list[dict]) -> dict:
+    if len(native) != len(reference) or not native:
+        raise ValueError("expected the same positive number of positions on both sides")
+    exceptions = []
+    for position, (a, b) in enumerate(zip(native, reference)):
+        for side, item in (("native", a), ("reference", b)):
+            if not all(math.isfinite(x) for x in item["logits"]):
+                raise ValueError(f"position {position}: non-finite {side} logits {item['logits']}")
+        actual, wanted = a["ids"][0], b["ids"][0]
+        if actual != wanted:
+            exceptions.append({"position": position, "expected": wanted, "actual": actual, **verdict(a, b, wanted, actual)})
+    return {"positions": len(native), "strict_equal": len(native) - len(exceptions), "exceptions": exceptions,
+            "policy_pass": all(item["verdict"] == "near-tie" for item in exceptions)}
+
+
+def load(path: Path, vocab: int):
+    import numpy as np
+
+    words = np.fromfile(path, dtype=np.uint16)
+    if words.size % vocab:
+        raise ValueError(f"{path}: not whole rows of {vocab} bf16 logits")
+    return (words.astype(np.uint32) << 16).view(np.float32).reshape(-1, vocab)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--reference", type=Path, required=True)
+    parser.add_argument("--native", type=Path, required=True)
+    parser.add_argument("--vocab", type=int, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    ref, nat = load(args.reference, args.vocab), load(args.native, args.vocab)
+    try:
+        sides = []
+        for name, rows in (("native", nat), ("reference", ref)):
+            records = []
+            for position, row in enumerate(rows):
+                try:
+                    records.append(top2(row))
+                except ValueError as err:
+                    raise ValueError(f"{name} position {position}: {err}") from None
+            sides.append(records)
+        result = judge(*sides)
+    except ValueError as err:
+        print(f"refused: {err}")
+        raise SystemExit(1)
+    result["sha256"] = {key: hashlib.sha256(path.read_bytes()).hexdigest() for key, path in (("reference", args.reference), ("native", args.native))}
+    with args.output.open("x") as out:
+        out.write(json.dumps(result, indent=2) + "\n")
+    for item in result["exceptions"]:
+        print(f"position {item['position']}: {item['verdict']}: reference {item['python_top2']}, native {item['native_top2']}")
+    print(f"{result['strict_equal']} of {result['positions']} positions equal; {len(result['exceptions'])} differ; "
+          f"policy {'pass' if result['policy_pass'] else 'FAIL'}")
+    if not result["policy_pass"]:
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

tools/zig/llama_teacher_judge_test.py

@@ -0,0 +1,113 @@
+"""The Llama teacher judge passes equal rows and one-step ties, and is red on wider margins and other contenders."""
+import copy
+import math
+import os
+import struct
+import subprocess
+import sys
+import tempfile
+import unittest
+
+from llama_teacher_judge import judge, top2
+
+try:
+    import numpy
+except ImportError:
+    numpy = None
+
+
+def rows(n=8):
+    return [{"ids": [1, 2], "logits": [38.25, 38.0]} for _ in range(n)]
+
+
+class LlamaTeacherJudgeTest(unittest.TestCase):
+    def test_equal(self):
+        result = judge(rows(), rows())
+        self.assertEqual(result["strict_equal"], 8)
+        self.assertTrue(result["policy_pass"])
+
+    def test_one_step_tie(self):
+        native = rows()
+        native[3] = {"ids": [2, 1], "logits": [38.0, 38.0]}
+        result = judge(native, rows())
+        self.assertTrue(result["policy_pass"])
+        self.assertEqual([item["position"] for item in result["exceptions"]], [3])
+
+    def test_wider_margin_is_red(self):
+        native = rows()
+        native[0] = {"ids": [2, 1], "logits": [38.5, 38.0]}
+        self.assertFalse(judge(native, rows())["policy_pass"])
+
+    def test_third_contender_is_red(self):
+        # three tokens within a step (#1415's position 1370): the two sides' contenders differ
+        native = rows()
+        native[5] = {"ids": [0, 1], "logits": [38.0, 38.0]}
+        reference = copy.deepcopy(rows())
+        reference[5] = {"ids": [1, 2], "logits": [38.25, 38.25]}
+        self.assertFalse(judge(native, reference)["policy_pass"])
+
+    def test_non_finite_rows_are_refused_on_either_side(self):
+        for bad in (math.nan, math.inf, -math.inf):
+            for side in ("native", "reference"):
+                native, reference = rows(), rows()
+                row = {"ids": [1, 2], "logits": [bad, 38.0]}
+                (native if side == "native" else reference)[4] = row
+                with self.assertRaises(ValueError, msg=f"{bad} on the {side} side"):
+                    judge(native, reference)
+
+    def test_non_finite_with_the_same_winner_is_refused(self):
+        # the winning id unchanged on both sides must not hide the score
+        native, reference = rows(), rows()
+        native[2] = {"ids": [1, 2], "logits": [math.nan, 0.0]}
+        reference[2] = {"ids": [1, 2], "logits": [math.nan, 0.0]}
+        with self.assertRaises(ValueError):
+            judge(native, reference)
+        native[2] = {"ids": [1, 2], "logits": [38.25, -math.inf]}
+        with self.assertRaises(ValueError):
+            judge(native, rows())
+
+    @unittest.skipUnless(numpy, "numpy")
+    def test_top2_refuses_a_non_finite_row_anywhere(self):
+        for bad in (numpy.nan, numpy.inf, -numpy.inf):
+            row = numpy.array([1.0, 3.0, 2.0, 0.5], dtype=numpy.float32)
+            row[3] = bad  # not the winner: the winning id would be unchanged
+            with self.assertRaises(ValueError):
+                top2(row)
+
+    @unittest.skipUnless(numpy, "numpy")
+    def test_the_checker_exits_1_on_a_non_finite_row(self):
+        def bf16(values):
+            return b"".join(struct.pack("<H", struct.unpack("<I", struct.pack("<f", v))[0] >> 16) for v in values)
+        good = [1.0, 3.0, 2.0, 0.5] * 2
+        here = os.path.dirname(os.path.abspath(__file__))
+        for side in ("native", "reference"):
+            for bad in (math.nan, math.inf, -math.inf):
+                with tempfile.TemporaryDirectory() as d:
+                    planted = list(good)
+                    planted[4 + 3] = bad  # row 1, not its winner
+                    files = {"native": good, "reference": good}
+                    files[side] = planted
+                    for name, values in files.items():
+                        with open(os.path.join(d, name + ".bin"), "wb") as f:
+                            f.write(bf16(values))
+                    run = subprocess.run([sys.executable, os.path.join(here, "llama_teacher_judge.py"), "--native", os.path.join(d, "native.bin"),
+                                          "--reference", os.path.join(d, "reference.bin"), "--vocab", "4", "--output", os.path.join(d, "out.json")],
+                                         capture_output=True, text=True, cwd=here)
+                    self.assertEqual(run.returncode, 1, f"{bad} on the {side} side: {run.stdout}{run.stderr}")
+                    self.assertIn("refused", run.stdout)
+        with tempfile.TemporaryDirectory() as d:  # the control: the same files, all finite, pass
+            for name in ("native", "reference"):
+                with open(os.path.join(d, name + ".bin"), "wb") as f:
+                    f.write(bf16(good))
+            run = subprocess.run([sys.executable, os.path.join(here, "llama_teacher_judge.py"), "--native", os.path.join(d, "native.bin"),
+                                  "--reference", os.path.join(d, "reference.bin"), "--vocab", "4", "--output", os.path.join(d, "out.json")],
+                                 capture_output=True, text=True, cwd=here)
+            self.assertEqual(run.returncode, 0, run.stdout + run.stderr)
+
+    def test_lengths_must_match(self):
+        with self.assertRaises(ValueError):
+            judge(rows(7), rows(8))
+
+
+if __name__ == "__main__":
+    unittest.main()

zig/build/metal_targets.zig

@@ -327,6 +327,7 @@ pub fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: s
         .{ .name = "tf-gemma4-bench", .path = "zig/tests/gemma4/bench.zig", .about = "Gemma 4 dense step costs: projections alone by rows against whole forwards" },
         .{ .name = "tf-gemma4-no-tensor", .path = "zig/tests/gemma4/no_tensor.zig", .about = "Gemma 4's Metal kernels on a Mac without tensor units: what builds, and the row path's bits" },
         .{ .name = "tf-spark-run", .path = "zig/tests/spark25/run.zig", .about = "Spark2.5 on Metal: greedy tokens, every position's logits or a prompt's last logits, for the reference's check" },
+        .{ .name = "tf-llama-run", .path = "zig/tests/llama/run.zig", .about = "Llama on Metal: greedy tokens, every position's logits or a prompt's last logits, for the reference's check" },
         .{ .name = "tf-qwen36-row-seam", .path = "zig/tests/qwen36_row_seam.zig", .about = "Model-free fixture: the row projection's bias sums stay fp32 at every width" },
     };
     for (engine_programs) |p| {