model/needle: tests — oracle parity, unit maths, tokenizer receipts

Proves the reference against a PyTorch oracle (modeling_needle.py + real
weights) and unit-checks every primitive:
- Generate reproduces the exact coherent call from the real weights, plus two
  further tool-use prompts (generalisation, not a memorised fixture).
- Exact encoder-input + generated token parity; encoder-hidden and decoder
  cross-attention-step numerics matched to the oracle; first-token top-5 logit
  parity (ids + magnitudes) — the gold receipt.
- ZCRMSNorm formula (incl. no-mean-centring guard), NeoX RoPE, GQA linear
  algebra, bf16 widen, and SentencePiece encode/decode against oracle ids.

Real-weight tests skip cleanly when the checkpoint is absent (CI stays green).

Co-Authored-By: Virgil <virgil@lethean.io>

Author: Snider

go/model/needle/config_test.go

@@ -0,0 +1,51 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import "testing"
+
+// TestConfig_DefaultConfig_Dims pins the published Needle geometry the reference
+// depends on: 512-d, 8 query / 4 kv heads (head_dim 64), 12 encoder / 8 decoder
+// layers, vocab 8192.
+func TestConfig_DefaultConfig_Dims(t *testing.T) {
+	c := DefaultConfig()
+	if c.HeadDim() != 64 {
+		t.Errorf("HeadDim = %d, want 64", c.HeadDim())
+	}
+	if c.KVDim() != 256 {
+		t.Errorf("KVDim = %d, want 256", c.KVDim())
+	}
+	if c.NumEncoderLayers != 12 || c.NumDecoderLayers != 8 {
+		t.Errorf("layers = %d/%d, want 12/8", c.NumEncoderLayers, c.NumDecoderLayers)
+	}
+	if c.EosTokenID != 1 || c.ToolsTokenID != 5 {
+		t.Errorf("special ids eos=%d tools=%d, want 1/5", c.EosTokenID, c.ToolsTokenID)
+	}
+}
+
+// TestConfig_LoadConfig_RealCheckpoint reads the real config.json and confirms it
+// matches the published geometry (and that a present file overrides defaults).
+func TestConfig_LoadConfig_RealCheckpoint(t *testing.T) {
+	cfg, err := LoadConfig(snapshotDir)
+	if err != nil {
+		t.Fatalf("LoadConfig: %v", err)
+	}
+	if cfg.HiddenSize != 512 || cfg.VocabSize != 8192 {
+		t.Errorf("hidden/vocab = %d/%d, want 512/8192", cfg.HiddenSize, cfg.VocabSize)
+	}
+	if cfg.NumKVHeads != 4 || cfg.NumHeads != 8 {
+		t.Errorf("heads = %d/%d, want 8/4", cfg.NumHeads, cfg.NumKVHeads)
+	}
+}
+
+// TestConfig_LoadConfig_MissingIsDefault confirms a directory without config.json
+// yields DefaultConfig rather than an error — a bare weights+tokenizer pack runs.
+func TestConfig_LoadConfig_MissingIsDefault(t *testing.T) {
+	cfg, err := LoadConfig(t.TempDir())
+	if err != nil {
+		t.Fatalf("LoadConfig(empty dir): %v", err)
+	}
+	if cfg != DefaultConfig() {
+		t.Errorf("missing config.json did not fall back to DefaultConfig")
+	}
+}

go/model/needle/generate_parity_test.go

@@ -0,0 +1,80 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import "testing"
+
+// TestGenerate_Generalises runs the reference on tool-use prompts beyond the
+// primary fixture, each checked against its PyTorch-oracle output. Matching more
+// than one prompt shows the port is the architecture, not a memorised fixture.
+func TestGenerate_Generalises(t *testing.T) {
+	m := loadTestModel(t)
+	cases := []struct {
+		name, query, tools, want string
+	}{
+		{
+			name:  "send_email",
+			query: "Send an email to john@example.com saying hello",
+			tools: `[{"name":"send_email","description":"Send an email to a recipient.","parameters":{"to":{"type":"string","description":"The recipient email address.","required":true},"body":{"type":"string","description":"The email body text.","required":true}}}]`,
+			want:  ` [{"name":"send_email","arguments":{"to":"john@example.com","body":"hello"}}]`,
+		},
+		{
+			name:  "get_stock_price",
+			query: "Get the current stock price of AAPL",
+			tools: `[{"name":"get_stock_price","description":"Get the current stock price.","parameters":{"symbol":{"type":"string","description":"Ticker symbol.","required":true}}}]`,
+			want:  ` [{"name":"get_stock_price","arguments":{"symbol":"AAPL"}}]`,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := m.Generate(tc.query, tc.tools, 64); got != tc.want {
+				t.Fatalf("Generate(%s):\n got: %q\nwant: %q", tc.name, got, tc.want)
+			}
+		})
+	}
+}
+
+// TestGenerate_FirstTokenLogitParity is the gold oracle receipt: the reference's
+// first-step logits (decoder over [EOS]) must match the PyTorch oracle in both
+// the top-5 token ranking and the raw logit magnitudes (within f32-vs-bf16
+// tolerance), not merely the argmax.
+func TestGenerate_FirstTokenLogitParity(t *testing.T) {
+	m := loadTestModel(t)
+	encoderHidden := m.encode(oracleEncIDs)
+	dh := m.decode([]int{m.cfg.EosTokenID}, encoderHidden, len(oracleEncIDs))
+	logits := m.logits(dh[:m.cfg.HiddenSize])
+
+	wantIDs := []int{4, 294, 8063, 302, 264}
+	wantVals := []float32{1.22699, -14.80001, -16.13503, -17.17564, -17.1861}
+
+	gotIDs, gotVals := top5(logits)
+	for i := range wantIDs {
+		if gotIDs[i] != wantIDs[i] {
+			t.Errorf("top5 id[%d] = %d, want %d (full got %v)", i, gotIDs[i], wantIDs[i], gotIDs)
+		}
+		if !close32(gotVals[i], wantVals[i], 0.05) {
+			t.Errorf("top5 logit[%d] = %.5f, want %.5f", i, gotVals[i], wantVals[i])
+		}
+	}
+}
+
+// top5 returns the five highest-scoring ids and their logits, descending.
+func top5(v []float32) ([]int, []float32) {
+	ids := make([]int, 5)
+	vals := make([]float32, 5)
+	for i := range 5 {
+		vals[i] = -1e30
+	}
+	for id, x := range v {
+		for r := range 5 {
+			if x > vals[r] {
+				copy(ids[r+1:], ids[r:4])
+				copy(vals[r+1:], vals[r:4])
+				ids[r] = id
+				vals[r] = x
+				break
+			}
+		}
+	}
+	return ids, vals
+}

go/model/needle/linalg_test.go

@@ -0,0 +1,61 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import (
+	"math"
+	"testing"
+)
+
+// TestLinalg_linearNoBias_RowDot confirms y = x . Wᵀ: output o is x dotted with
+// weight row o (row-major [outDim, inDim]).
+func TestLinalg_linearNoBias_RowDot(t *testing.T) {
+	x := []float32{1, 2, 3}
+	// W = [[1,0,0],[0,1,0],[1,1,1]] flattened row-major.
+	w := []float32{1, 0, 0, 0, 1, 0, 1, 1, 1}
+	got := linearNoBias(x, w, 3, 3)
+	want := []float32{1, 2, 6}
+	for i := range want {
+		if got[i] != want[i] {
+			t.Errorf("linearNoBias[%d] = %v, want %v", i, got[i], want[i])
+		}
+	}
+}
+
+// TestLinalg_softmaxInPlace_Distribution confirms softmax sums to 1 and is
+// monotonic in the inputs.
+func TestLinalg_softmaxInPlace_Distribution(t *testing.T) {
+	s := []float32{1, 2, 3}
+	softmaxInPlace(s)
+	var sum float32
+	for _, v := range s {
+		sum += v
+	}
+	if math.Abs(float64(sum-1)) > 1e-6 {
+		t.Errorf("softmax sum = %v, want 1", sum)
+	}
+	if !(s[0] < s[1] && s[1] < s[2]) {
+		t.Errorf("softmax not monotonic: %v", s)
+	}
+}
+
+// TestLinalg_sigmoid_Midpoint confirms sigmoid(0) = 0.5, the init value of every
+// residual gate (weight init 0 -> half-open gate).
+func TestLinalg_sigmoid_Midpoint(t *testing.T) {
+	if got := sigmoid(0); math.Abs(float64(got-0.5)) > 1e-6 {
+		t.Errorf("sigmoid(0) = %v, want 0.5", got)
+	}
+}
+
+// TestLinalg_clipAdd_Clamps confirms the bf16-range clamp mirrors _add_clipped.
+func TestLinalg_clipAdd_Clamps(t *testing.T) {
+	if got := clipAdd(65000, 1000); got != 65500 {
+		t.Errorf("clipAdd over-range = %v, want 65500", got)
+	}
+	if got := clipAdd(-65000, -1000); got != -65500 {
+		t.Errorf("clipAdd under-range = %v, want -65500", got)
+	}
+	if got := clipAdd(1.5, 2.5); got != 4 {
+		t.Errorf("clipAdd in-range = %v, want 4", got)
+	}
+}

go/model/needle/model_test.go

@@ -0,0 +1,117 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import (
+	"math"
+	"slices"
+	"testing"
+
+	coreio "dappco.re/go/io"
+)
+
+// snapshotDir is the local Needle checkpoint. Tests that need the real weights
+// skip when it is absent (e.g. CI), so the suite stays green without the model
+// while giving a real end-to-end receipt on a machine that has it.
+const snapshotDir = "/Users/snider/.cache/huggingface/hub/models--Cactus-Compute--needle/snapshots/5f89b4307696d669c3df1d38ae057e6e1728b107"
+
+// oracleQuery/oracleTools are the exact prompt the PyTorch reference was run on;
+// the fixtures below are that run's output (see the lane report).
+const (
+	oracleQuery = "What is the weather in San Francisco?"
+	oracleTools = `[{"name":"get_weather","description":"Get current weather for a city.","parameters":{"location":{"type":"string","description":"City name.","required":true}}}]`
+)
+
+var (
+	oracleEncIDs = []int{4279, 743, 302, 1149, 362, 711, 327, 1295, 1075, 378, 275, 8047, 8105, 5, 356, 294, 264, 358, 8062, 1331, 265, 283, 264, 618, 407, 1149, 345, 289, 2082, 284, 318, 282, 506, 282, 298, 264, 315, 265, 283, 264, 2523, 417, 284, 301, 262, 312, 434}
+	oracleGenIDs = []int{4, 356, 294, 264, 358, 8062, 1331, 265, 393, 282, 506, 264, 8074, 327, 1295, 1075, 378, 275, 8047, 503}
+	oracleOutput = ` [{"name":"get_weather","arguments":{"location":"San Francisco"}}]`
+)
+
+func loadTestModel(t *testing.T) *Model {
+	t.Helper()
+	if !coreio.Local.Exists(snapshotDir + "/model.safetensors") {
+		t.Skip("needle checkpoint not present at " + snapshotDir)
+	}
+	m, err := Load(snapshotDir)
+	if err != nil {
+		t.Fatalf("Load: %v", err)
+	}
+	return m
+}
+
+func close32(a, b, tol float32) bool { return float32(math.Abs(float64(a-b))) <= tol }
+
+// TestModel_Generate_ToolCall is the headline receipt: from the real weights, the
+// reference greedily generates the exact coherent function call the PyTorch
+// oracle produced.
+func TestModel_Generate_ToolCall(t *testing.T) {
+	m := loadTestModel(t)
+	got := m.Generate(oracleQuery, oracleTools, 64)
+	if got != oracleOutput {
+		t.Fatalf("Generate mismatch:\n got: %q\nwant: %q", got, oracleOutput)
+	}
+	t.Logf("VERBATIM GENERATION: %s", got)
+}
+
+// TestModel_generateIDs_TokenParity pins exact encoder-input and generated token
+// ids against the oracle — a stricter check than the decoded string.
+func TestModel_generateIDs_TokenParity(t *testing.T) {
+	m := loadTestModel(t)
+	enc, gen := m.generateIDs(oracleQuery, oracleTools, 64)
+	if !slices.Equal(enc, oracleEncIDs) {
+		t.Fatalf("encoder ids mismatch:\n got: %v\nwant: %v", enc, oracleEncIDs)
+	}
+	if !slices.Equal(gen, oracleGenIDs) {
+		t.Fatalf("generated ids mismatch:\n got: %v\nwant: %v", gen, oracleGenIDs)
+	}
+}
+
+// TestModel_encode_EncoderHidden checks the encoder pass numerically against the
+// oracle's encoder output (first token, last token, and the whole-tensor sum),
+// isolating encoder correctness from the decoder.
+func TestModel_encode_EncoderHidden(t *testing.T) {
+	m := loadTestModel(t)
+	h := m.encode(oracleEncIDs)
+	if len(h) != len(oracleEncIDs)*m.cfg.HiddenSize {
+		t.Fatalf("encoder hidden len = %d, want %d", len(h), len(oracleEncIDs)*m.cfg.HiddenSize)
+	}
+	wantT0 := []float32{-1.44818, 0.02975, 0.43704, 0.51372, -1.47645}
+	wantTlast := []float32{-0.62455, 0.66252, 0.34196, 0.43071, 0.47612}
+	hidden := m.cfg.HiddenSize
+	last := (len(oracleEncIDs) - 1) * hidden
+	for i := range wantT0 {
+		if !close32(h[i], wantT0[i], 0.02) {
+			t.Errorf("enc_hid[0][%d] = %.5f, want %.5f", i, h[i], wantT0[i])
+		}
+		if !close32(h[last+i], wantTlast[i], 0.02) {
+			t.Errorf("enc_hid[last][%d] = %.5f, want %.5f", i, h[last+i], wantTlast[i])
+		}
+	}
+	var sum float32
+	for _, v := range h {
+		sum += v
+	}
+	if !close32(sum, -730.317, 2.0) {
+		t.Errorf("encoder hidden sum = %.3f, want ~-730.317", sum)
+	}
+}
+
+// TestModel_decode_CrossAttentionStep checks the decoder-with-cross-attention
+// step: the decoder hidden (post-norm) after one step over [EOS], against the
+// oracle. This exercises masked self-attn + cross-attn + gates end to end.
+func TestModel_decode_CrossAttentionStep(t *testing.T) {
+	m := loadTestModel(t)
+	encoderHidden := m.encode(oracleEncIDs)
+	dh := m.decode([]int{m.cfg.EosTokenID}, encoderHidden, len(oracleEncIDs))
+	want := []float32{-0.01168, 0.01785, -0.14802, -0.12168, 0.00175}
+	for i := range want {
+		if !close32(dh[i], want[i], 0.02) {
+			t.Errorf("dec_hid[step0][%d] = %.5f, want %.5f", i, dh[i], want[i])
+		}
+	}
+	// First-token argmax must be <tool_call> (id 4).
+	if got := argmax(m.logits(dh[:m.cfg.HiddenSize])); got != 4 {
+		t.Errorf("first-token argmax = %d, want 4 (<tool_call>)", got)
+	}
+}

go/model/needle/norm_test.go

@@ -0,0 +1,51 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import (
+	"math"
+	"testing"
+)
+
+// TestNorm_zcRMSNorm_Formula pins the exact ZCRMSNorm maths: variance is the mean
+// of squares (NOT centred), and the weight enters as (1 + weight). x = [1,2,3,4],
+// weight = 0 -> x / rms(x), rms = sqrt(30/4).
+func TestNorm_zcRMSNorm_Formula(t *testing.T) {
+	x := []float32{1, 2, 3, 4}
+	weight := []float32{0, 0, 0, 0}
+	inv := float32(1.0 / math.Sqrt(7.5+1e-6))
+	want := []float32{1 * inv, 2 * inv, 3 * inv, 4 * inv}
+	got := zcRMSNorm(x, weight, 1e-6)
+	for i := range want {
+		if math.Abs(float64(got[i]-want[i])) > 1e-5 {
+			t.Errorf("zcRMSNorm[%d] = %.6f, want %.6f", i, got[i], want[i])
+		}
+	}
+}
+
+// TestNorm_zcRMSNorm_WeightShift confirms the (1 + weight) scale: a weight of 1
+// doubles the normalised output, a weight of 0 leaves it, matching the
+// zero-centred convention (init-0 weight => identity scale).
+func TestNorm_zcRMSNorm_WeightShift(t *testing.T) {
+	x := []float32{2, 0, 0, 0} // rms = sqrt(4/4) = 1, so normalised == x
+	base := zcRMSNorm(x, []float32{0, 0, 0, 0}, 1e-6)
+	if math.Abs(float64(base[0]-2)) > 1e-4 {
+		t.Fatalf("weight 0: got %.6f, want 2", base[0])
+	}
+	shifted := zcRMSNorm(x, []float32{1, 0, 0, 0}, 1e-6)
+	if math.Abs(float64(shifted[0]-4)) > 1e-4 {
+		t.Fatalf("weight 1: got %.6f, want 4 (2 * (1+1))", shifted[0])
+	}
+}
+
+// TestNorm_zcRMSNorm_NotMeanCentred guards the classic mistake: if the mean were
+// subtracted, a constant vector would normalise to zero. ZCRMSNorm must NOT do
+// that — a constant vector normalises to +-1 * (1+weight).
+func TestNorm_zcRMSNorm_NotMeanCentred(t *testing.T) {
+	got := zcRMSNorm([]float32{3, 3, 3, 3}, []float32{0, 0, 0, 0}, 1e-6)
+	for i, v := range got {
+		if math.Abs(float64(v-1)) > 1e-4 {
+			t.Errorf("constant-vector norm[%d] = %.6f, want 1 (mean must NOT be subtracted)", i, v)
+		}
+	}
+}