llama: published MLX affine quants (8-bit and 4-bit) read in place, beside the bf16 default (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,86 @@
+"""A Llama checkpoint's teacher-forced logits in float64 on the CPU, rounded once to bf16: the exact reference."""
+# usage: python3 tools/zig/llama_f64_reference.py MODEL_DIR tokens.json out.bin (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) -> None:
+    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):
+        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)
+        return scales * q + biases
+
+    def norm(x, g):
+        return g * (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()
+        x = w("model.embed_tokens.weight")[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 = (h @ w(p + "self_attn.q_proj.weight").T).view(n, heads, dim).transpose(0, 1)
+            k = (h @ w(p + "self_attn.k_proj.weight").T).view(n, kv, dim).transpose(0, 1)
+            v = (h @ w(p + "self_attn.v_proj.weight").T).view(n, kv, dim).transpose(0, 1)
+            q, k = q * cos + rotate(q) * sin, k * cos + rotate(k) * sin
+            k, v = k.repeat_interleave(heads // kv, 0), v.repeat_interleave(heads // kv, 0)
+            a = torch.softmax(q @ k.transpose(1, 2) / dim**0.5 + mask, -1) @ v
+            x = x + 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"))
+            x = x + (torch.nn.functional.silu(h @ w(p + "mlp.gate_proj.weight").T) * (h @ w(p + "mlp.up_proj.weight").T)) @ 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])

tools/zig/llama_teacher_judge.py

@@ -0,0 +1,65 @@
+"""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
+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 (the engine's argmax and numpy's)."""
+    import numpy as np
+
+    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)):
+        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)
+    result = judge([top2(row) for row in nat], [top2(row) for row in ref])
+    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,44 @@
+"""The Llama teacher judge passes equal rows and one-step ties, and is red on wider margins and other contenders."""
+import copy
+import unittest
+
+from llama_teacher_judge import judge
+
+
+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_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| {