model/needle: pure-Go host-f32 reference forward for the Needle encoder-decoder

Needle is Cactus-Compute's 26M-param encoder-decoder ("Simple Attention
Network") for on-device function calling. This is a self-contained, pure-Go
host reference (no engine, no cgo, no GPU) that loads the published safetensors
weights, widens bf16 to f32, and runs the full encoder + cross-attending
decoder to de-risk the encoder-decoder + cross-attention direction before any
accelerated port.

Ported field-for-field from the cactus modeling_needle.py:
- ZCRMSNorm: RMSNorm with no mean-centring, weight enters as (1 + weight).
- Gated residuals: clip(h + sigmoid(gate_scalar) * attn_branch), pre-norm.
- GQA (8 query / 4 kv heads) with per-head QK-norm before RoPE (self-attn only).
- Decoder cross-attends to the FINAL encoder output, own k/v per layer, no RoPE.
- Shared embed_tokens scaled by sqrt(d); tied unscaled lm_head.
- Minimal SentencePiece BPE encode/decode (byte-fallback) via the shared reader.

Reuses model/safetensors (weights) and decode/tokenizer (SPM proto reader);
all forward maths is local f32. Registration and the engine port are later
slices.

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

Author: Snider

go/model/needle/attention.go

@@ -0,0 +1,97 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import "math"
+
+// attention is the one attention block Needle reuses for encoder self-attention,
+// decoder masked self-attention and decoder cross-attention (NeedleAttention).
+// It is grouped-query (8 query heads share 4 key/value heads, kvHead = h/group),
+// applies per-head QK-norm before RoPE, and scales scores by 1/sqrt(head_dim).
+//
+//   - prefix names the weight group, e.g. "model.decoder.layers.0.encoder_attn".
+//   - xq/xkv are already-normalised hidden states ([qLen,hidden] / [kvLen,hidden]);
+//     for self-attention they are the same slice.
+//   - rope is nil for cross-attention (no positional rotation there).
+//   - causal masks key j>i (decoder self-attention); false is bidirectional.
+//
+// Returns the out_proj'd result [qLen, hidden].
+func (m *Model) attention(prefix string, xq []float32, qLen int, xkv []float32, kvLen int, rope *ropeTable, causal bool) []float32 {
+	c := m.cfg
+	hidden := c.HiddenSize
+	headDim := c.HeadDim()
+	kvDim := c.KVDim()
+	group := c.NumHeads / c.NumKVHeads
+	scale := float32(1.0 / math.Sqrt(float64(headDim)))
+
+	qProj := m.w.get(prefix + ".q_proj.weight")
+	kProj := m.w.get(prefix + ".k_proj.weight")
+	vProj := m.w.get(prefix + ".v_proj.weight")
+	oProj := m.w.get(prefix + ".out_proj.weight")
+	qNorm := m.w.get(prefix + ".q_norm.weight")
+	kNorm := m.w.get(prefix + ".k_norm.weight")
+
+	// Project then reshape to [len, heads, head_dim] (flattened row-major).
+	q := linearRows(xq, qLen, qProj, hidden, hidden)  // [qLen, numHeads*headDim]
+	k := linearRows(xkv, kvLen, kProj, kvDim, hidden) // [kvLen, numKV*headDim]
+	v := linearRows(xkv, kvLen, vProj, kvDim, hidden)
+
+	// Per-head QK-norm, then RoPE (self-attention only).
+	for i := range qLen {
+		for h := range c.NumHeads {
+			head := q[i*hidden+h*headDim : i*hidden+h*headDim+headDim]
+			copy(head, zcRMSNorm(head, qNorm, c.RMSNormEps))
+			if rope != nil {
+				rope.apply(head, i)
+			}
+		}
+	}
+	for j := range kvLen {
+		for h := range c.NumKVHeads {
+			head := k[j*kvDim+h*headDim : j*kvDim+h*headDim+headDim]
+			copy(head, zcRMSNorm(head, kNorm, c.RMSNormEps))
+			if rope != nil {
+				rope.apply(head, j)
+			}
+		}
+	}
+
+	// Attention per query head against its grouped kv head.
+	attnOut := make([]float32, qLen*hidden)
+	scores := make([]float32, kvLen)
+	for h := range c.NumHeads {
+		kvHead := h / group
+		for i := range qLen {
+			qv := q[i*hidden+h*headDim : i*hidden+h*headDim+headDim]
+			limit := kvLen
+			if causal {
+				limit = i + 1 // key j > i is masked out
+			}
+			for j := range kvLen {
+				if j >= limit {
+					scores[j] = float32(math.Inf(-1))
+					continue
+				}
+				kv := k[j*kvDim+kvHead*headDim : j*kvDim+kvHead*headDim+headDim]
+				var dot float32
+				for d := range headDim {
+					dot += qv[d] * kv[d]
+				}
+				scores[j] = dot * scale
+			}
+			softmaxInPlace(scores[:kvLen])
+			out := attnOut[i*hidden+h*headDim : i*hidden+h*headDim+headDim]
+			for j := range kvLen {
+				p := scores[j]
+				if p == 0 {
+					continue
+				}
+				vv := v[j*kvDim+kvHead*headDim : j*kvDim+kvHead*headDim+headDim]
+				for d := range headDim {
+					out[d] += p * vv[d]
+				}
+			}
+		}
+	}
+	return linearRows(attnOut, qLen, oProj, hidden, hidden)
+}

go/model/needle/config.go

@@ -0,0 +1,152 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+// Package needle is a pure-Go, host-side f32 reference forward for Cactus
+// Compute's "Needle" model — a 26M-parameter encoder-decoder ("Simple Attention
+// Network") distilled for on-device function calling. It has no engine, no cgo
+// and no GPU: it loads the published safetensors weights, widens bf16 to f32,
+// and runs the encoder + cross-attending decoder in plain Go. Its purpose is to
+// de-risk the encoder-decoder + cross-attention direction before any accelerated
+// port — a readable oracle whose only claim to correctness is that it reproduces
+// the reference model's tokens.
+//
+//	m, err := needle.Load("/path/to/needle-snapshot")
+//	if err != nil {
+//		return err
+//	}
+//	out := m.Generate("What is the weather in San Francisco?",
+//		`[{"name":"get_weather","parameters":{"location":"string"}}]`, 64)
+//	// out == ` [{"name":"get_weather","arguments":{"location":"San Francisco"}}]`
+//
+// The architecture (228 weights, verified against the safetensors header):
+//   - Shared embed_tokens[vocab,d] scaled by sqrt(d); tied lm_head.
+//   - Encoder x12: input_layernorm -> GQA self-attn (per-head QK-norm, RoPE,
+//     bidirectional) -> scalar sigmoid gate on the residual. No FFN. final_norm.
+//   - Decoder x8: input_layernorm -> masked self-attn (gate) ->
+//     encoder_attn_layer_norm -> cross-attn to the FINAL encoder output (own
+//     k/v, no RoPE, gate). No FFN. norm.
+package needle
+
+import (
+	core "dappco.re/go"
+	coreio "dappco.re/go/io"
+)
+
+// Config holds the Needle hyper-parameters read from a checkpoint's config.json.
+// The zero value is not usable; construct via LoadConfig or DefaultConfig.
+type Config struct {
+	VocabSize        int     // token count (8192)
+	HiddenSize       int     // d_model (512)
+	NumHeads         int     // query heads (8)
+	NumKVHeads       int     // key/value heads for GQA (4)
+	NumEncoderLayers int     // encoder depth (12)
+	NumDecoderLayers int     // decoder depth (8)
+	RopeTheta        float64 // RoPE base (10000)
+	RMSNormEps       float64 // ZCRMSNorm epsilon (1e-6)
+	PadTokenID       int     // 0
+	EosTokenID       int     // 1 — also the decoder start token
+	BosTokenID       int     // 2
+	ToolsTokenID     int     // 5 — <tools> separator between query and tools
+}
+
+// HeadDim is the per-head width (hidden / heads = 64).
+//
+//	cfg.HeadDim() // 64
+func (c Config) HeadDim() int { return c.HiddenSize / c.NumHeads }
+
+// KVDim is the concatenated key/value projection width (kv_heads * head_dim = 256).
+//
+//	cfg.KVDim() // 256
+func (c Config) KVDim() int { return c.NumKVHeads * c.HeadDim() }
+
+// DefaultConfig returns the published Needle configuration. It is the fallback
+// when a checkpoint ships no config.json, and the source of truth for the token
+// ids the tokenizer and decoder rely on.
+//
+//	cfg := needle.DefaultConfig()
+//	cfg.HiddenSize // 512
+func DefaultConfig() Config {
+	return Config{
+		VocabSize:        8192,
+		HiddenSize:       512,
+		NumHeads:         8,
+		NumKVHeads:       4,
+		NumEncoderLayers: 12,
+		NumDecoderLayers: 8,
+		RopeTheta:        10000,
+		RMSNormEps:       1e-6,
+		PadTokenID:       0,
+		EosTokenID:       1,
+		BosTokenID:       2,
+		ToolsTokenID:     5,
+	}
+}
+
+// configJSON mirrors the subset of config.json this reference needs. Fields the
+// checkpoint omits keep their DefaultConfig value.
+type configJSON struct {
+	VocabSize        *int     `json:"vocab_size"`
+	HiddenSize       *int     `json:"hidden_size"`
+	NumHeads         *int     `json:"num_heads"`
+	NumKVHeads       *int     `json:"num_kv_heads"`
+	NumEncoderLayers *int     `json:"num_encoder_layers"`
+	NumDecoderLayers *int     `json:"num_decoder_layers"`
+	RopeTheta        *float64 `json:"rope_theta"`
+	RMSNormEps       *float64 `json:"rms_norm_eps"`
+	PadTokenID       *int     `json:"pad_token_id"`
+	EosTokenID       *int     `json:"eos_token_id"`
+	BosTokenID       *int     `json:"bos_token_id"`
+}
+
+// LoadConfig reads config.json from a checkpoint directory, falling back to
+// DefaultConfig for any field the file omits. A missing config.json is not an
+// error — the published defaults are returned — so the reference runs against a
+// bare weights+tokenizer pack.
+//
+//	cfg, err := needle.LoadConfig("/models/needle")
+func LoadConfig(dir string) (Config, error) {
+	cfg := DefaultConfig()
+	path := dir + "/config.json"
+	raw, err := coreio.Local.Read(path)
+	if err != nil {
+		// A checkpoint without config.json is served by the published defaults.
+		return cfg, nil
+	}
+	var j configJSON
+	if r := core.JSONUnmarshal(core.AsBytes(raw), &j); !r.OK {
+		return cfg, core.E("needle.LoadConfig", "parse "+path, r.Err())
+	}
+	if j.VocabSize != nil {
+		cfg.VocabSize = *j.VocabSize
+	}
+	if j.HiddenSize != nil {
+		cfg.HiddenSize = *j.HiddenSize
+	}
+	if j.NumHeads != nil {
+		cfg.NumHeads = *j.NumHeads
+	}
+	if j.NumKVHeads != nil {
+		cfg.NumKVHeads = *j.NumKVHeads
+	}
+	if j.NumEncoderLayers != nil {
+		cfg.NumEncoderLayers = *j.NumEncoderLayers
+	}
+	if j.NumDecoderLayers != nil {
+		cfg.NumDecoderLayers = *j.NumDecoderLayers
+	}
+	if j.RopeTheta != nil {
+		cfg.RopeTheta = *j.RopeTheta
+	}
+	if j.RMSNormEps != nil {
+		cfg.RMSNormEps = *j.RMSNormEps
+	}
+	if j.PadTokenID != nil {
+		cfg.PadTokenID = *j.PadTokenID
+	}
+	if j.EosTokenID != nil {
+		cfg.EosTokenID = *j.EosTokenID
+	}
+	if j.BosTokenID != nil {
+		cfg.BosTokenID = *j.BosTokenID
+	}
+	return cfg, nil
+}

go/model/needle/decoder.go

@@ -0,0 +1,45 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import core "dappco.re/go"
+
+// decode runs the full decoder over the current decoder token ids, cross-attending
+// to the encoder output. Each of the 8 layers is masked self-attention then
+// cross-attention, both pre-norm with their own scalar sigmoid gate, no FFN:
+//
+//	normed = input_layernorm(h)
+//	h      = clip(h + sigmoid(self_attn_gate)  * self_attn(normed, normed))   # causal, RoPE
+//	cross  = encoder_attn(encoder_attn_layer_norm(h), encoderHidden)          # cross, NO RoPE
+//	h      = clip(h + sigmoid(cross_attn_gate) * cross)
+//
+// The cross-attention keys/values come from the FINAL encoder output (the same
+// encoderHidden for every layer, each with its own k/v projection). Returns the
+// decoder hidden states [decLen, hidden] after the trailing norm.
+func (m *Model) decode(decIDs []int, encoderHidden []float32, encLen int) []float32 {
+	c := m.cfg
+	hidden := c.HiddenSize
+	decLen := len(decIDs)
+
+	h := m.embed(decIDs) // [decLen, hidden], scaled
+	rope := newRopeTable(c.HeadDim(), c.RopeTheta, decLen)
+
+	for layer := range c.NumDecoderLayers {
+		p := core.Sprintf("model.decoder.layers.%d", layer)
+
+		normed := zcRMSNormRows(h, decLen, hidden, m.w.get(p+".input_layernorm.weight"), c.RMSNormEps)
+		selfAttn := m.attention(p+".self_attn", normed, decLen, normed, decLen, rope, true)
+		selfGate := sigmoid(m.w.get(p + ".self_attn_gate")[0])
+		for i := range decLen * hidden {
+			h[i] = clipAdd(h[i], selfGate*selfAttn[i])
+		}
+
+		crossNormed := zcRMSNormRows(h, decLen, hidden, m.w.get(p+".encoder_attn_layer_norm.weight"), c.RMSNormEps)
+		cross := m.attention(p+".encoder_attn", crossNormed, decLen, encoderHidden, encLen, nil, false)
+		crossGate := sigmoid(m.w.get(p + ".cross_attn_gate")[0])
+		for i := range decLen * hidden {
+			h[i] = clipAdd(h[i], crossGate*cross[i])
+		}
+	}
+	return zcRMSNormRows(h, decLen, hidden, m.w.get("model.decoder.norm.weight"), c.RMSNormEps)
+}

go/model/needle/encoder.go

@@ -0,0 +1,34 @@
+// SPDX-Licence-Identifier: EUPL-1.2
+
+package needle
+
+import core "dappco.re/go"
+
+// encode runs the full bidirectional encoder over token ids: embed (scaled by
+// sqrt(hidden)), then 12 gated self-attention layers with no FFN, then final_norm.
+// Each layer is pre-norm with a scalar sigmoid gate on the attention branch:
+//
+//	normed = input_layernorm(h)
+//	attn   = self_attn(normed, normed)           # bidirectional, RoPE
+//	h      = clip(h + sigmoid(attn_gate) * attn) # no FFN
+//
+// Returns the encoder hidden states [seqLen, hidden] the decoder cross-attends to.
+func (m *Model) encode(ids []int) []float32 {
+	c := m.cfg
+	hidden := c.HiddenSize
+	seqLen := len(ids)
+
+	h := m.embed(ids) // [seqLen, hidden], already scaled
+	rope := newRopeTable(c.HeadDim(), c.RopeTheta, seqLen)
+
+	for layer := range c.NumEncoderLayers {
+		p := core.Sprintf("model.encoder.layers.%d", layer)
+		normed := zcRMSNormRows(h, seqLen, hidden, m.w.get(p+".input_layernorm.weight"), c.RMSNormEps)
+		attn := m.attention(p+".self_attn", normed, seqLen, normed, seqLen, rope, false)
+		gate := sigmoid(m.w.get(p + ".attn_gate")[0])
+		for i := range seqLen * hidden {
+			h[i] = clipAdd(h[i], gate*attn[i])
+		}
+	}
+	return zcRMSNormRows(h, seqLen, hidden, m.w.get("model.encoder.final_norm.weight"), c.RMSNormEps)
+}