Back to blog
Guides

Porting Needle 3

Notes for anyone writing their own Needle 3 runtime, from the first community port: which oracle to test against and how to read the numbers, the tensor order that the .cact container promises, what the prompt looks like on the wire, how the ladder picks blocks, and how retrieval works while the release ships without a contrastive head.

HN

Henry Ndubuaku

||10 min read

The first community port of Needle 3 landed within days of the release, written in Rust against nothing but the .cact file and the Python reference. Its author sent back three things that cost him debugging time and would have cost the next person the same. This post is those three, written down properly, plus the parts of the runtime contract a port needs that no other guide states in one place: the prompt on the wire, the ladder rule, and the format's guarantees.

What you have to work from is public. The engines for thirteen platforms are on Hugging Face beside the weights, needle3.cact at CQ2 and checkpoints/needle3.safetensors as the half-precision master. The reference model is needle.model.architecture.SimpleAttentionNetwork in the cactus-needle package, and the container is documented in needle/model/export.py and in the format post. The architecture itself is in Simple Attention Networks and the Hadamard MLP.

1. Test against float32 with quantisation simulated, and read cosines, not relative error

The reference checkpoint carries dtype: bfloat16 in its config, so the obvious oracle is the JAX model exactly as it loads. Measure a correct port against that oracle with per-element relative error and the number you get is enormous, because relative error on a logit near zero is meaningless and most of an 8192-way logit vector sits near zero. The port is not wrong; the metric is.

Two things are being conflated. bfloat16 keeps 8 bits of mantissa, about two and a half decimal digits, and after twenty layers the reference's own logits carry that noise. On the prompt below, the reference in bfloat16 against itself in float32 agrees to a cosine of 0.999999 with a median relative error of 0.5%, and the largest single logit moves by 0.3. That is the floor no port can beat, and it already produces sign flips on logits that round to zero.

The bigger gap is that needle3.cact does not hold the master weights. It holds them at 2.125 bits per weight (4.125 for the embedding and the lane maps), with int8 activations and an int8 KV cache, and the correct oracle for a port of the file is the reference run in float32 with that quantisation simulated. The package has both halves:

import dataclasses
import jax, jax.numpy as jnp, numpy as np
from needle.model.run import load_checkpoint
from needle.model.architecture import SimpleAttentionNetwork
from needle.model.quantize import cq_quantize, _is_quant_leaf
from needle.model.tokenizer import get_tokenizer

params, config = load_checkpoint("checkpoints/needle3.safetensors")
config = dataclasses.replace(config, dtype="float32")

def path_of(path):
    return "/".join(getattr(k, "key", getattr(k, "name", str(k))) for k in path)

def quantise(path, leaf):
    if not _is_quant_leaf(path, leaf):
        return leaf
    bits = 4 if path_of(path) == "embedding/embedding" or "mhc_phi" in path_of(path) else 2
    return cq_quantize(jnp.asarray(leaf), bits)

qparams = jax.tree_util.tree_map_with_path(quantise, params)
model = SimpleAttentionNetwork(config)
ids = [2] + get_tokenizer().encode(prompt_text)
logits = np.asarray(model.apply({"params": qparams}, jnp.array([ids]), quant=True))[0, -1]

cq_quantize rounds each group of 128 weights exactly as the exporter does: rotate by the Walsh-Hadamard matrix, split into an fp16 norm and a unit direction, snap the direction to the Lloyd-Max codebook. quant=True fake-quantises activations to int8 and the KV cache to int8 per head. That forward is what the engine implements.

The numbers to expect, on turn on the kitchen lights with one set_lights tool declared: the quantised float32 forward against the unquantised one has a cosine of 0.9995 and a median relative error of 9%, and the two admissible first tokens, <think> and <tool_call>, swap places. The shipped engine on the same 67 tokens ranks them the same way as the quantised reference and puts the rest of the vocabulary nine logits below. Digits will not match, because accumulation order differs; the ranking and the cosine do.

So the acceptance test for a port is three checks, in this order. The tokenizer reproduces the engine's ids byte for byte (the debug build prints prefix ids and turn ids; the tokenizer is a plain SentencePiece BPE whose pieces and scores are the RAW tensor at the end of the file, and RefTokenizer in export.py is the reference encoder). The logits agree with the quantised float32 reference to a cosine above 0.999 and the same top choice on every prompt in your set. The finished JSON is identical to needle_complete on the sandbox presets, which is the only comparison a user of your port will ever make.

2. The tensor order is the contract

The directory in a .cact file is positional and nameless: 44 bytes per record, dtype, shape, offset, size, group and bits, and nothing that says which weight it is. Per-tensor checks therefore cannot catch a reordering, and a port that pins the order by reading the exporter's private _tensors list is one refactor away from breaking silently. The right fix is to state the order as part of the format, so here it is. The order below is what tag 0x05E12A84 promises; a change to it is a new tag, and an engine handed the other tag refuses rather than guesses.

Position zero is the token embedding, CQ4, which is also the tied output head. Then, for each of the 20 blocks in order, 27 tensors:

#tensordtypeshape (20L, width 768)
1norm_inFP16768
2q_projCQ2576 × 768
3k_projCQ296 × 768
4v_projCQ2128 × 768
5q_tapsFP163 × 576
6k_tapsFP163 × 96
7v_tapsFP163 × 128
8q_normFP1648
9k_normFP1648
10gate_projCQ2768 × 768
11out_projCQ2768 × 768
12post_normFP16768
13attn_gateFP161
14pre_hadaFP16768
15 to 19d1, d2, b2, d3, d4FP161024 each
20 to 25w1a, w1b, w2a, w2b, w3a, w3bFP16Monarch factor pairs
26, 27cond_v, cond_uFP16rank-8 gate

After the blocks: the nine mHC tensors (a_pre, a_post, a_res, b_pre, b_post, b_res in FP16, then phi_pre, phi_post, phi_res as CQ4, each phi flattened to layers × lanes rows); the two Hadamard permutations hada_p1 and hada_p2 as FP32 index vectors of length 1024; for each engram site, in site order 3, 7, 11, 15, 19, tables (CQ2, num_tables × 18432 rows of 128), key_proj (CQ2), value_proj (CQ2), taps (FP16); final_norm (FP16); then the probe heads: a heads.manifest FP16 vector with one code per head (1 embedding, 2 confidence, 3 router), followed by each head's six tensors probes (CQ4), gain (FP16), query (CQ4), row_bias (FP16), proj (CQ4), bias (FP16); and last the RAW tokenizer. The shipped file has exactly one head, confidence, which is how 1 + 20 × 27 + 9 + 2 + 5 × 4 + 1 + 1 + 6 + 1 comes to 581 records.

Every matrix is stored pre-transposed as [out, in], so an output row is contiguous along the reduction axis and the quantisation groups run along it. Names in the header would cost 581 strings the engine never reads; the order costs nothing and is now written down. A port should assert the record count and the dtype pattern above at load time and refuse anything else.

3. Retrieval without a contrastive head

Needle 3's engine can retrieve tools: given a catalogue larger than five, it embeds every tool description once, embeds the query, keeps the top five by cosine, and renders only those into the prompt. That path switches on only when the archive carries an embedding head, and this release does not ship one, so on needle3.cact every declared tool goes into the prefix and retrieve_tools is absent from the Python package rather than empty. If your port sees no head with code 1 in the manifest, retrieval is off, and the file is the source of truth for that decision.

What the release does ship is needle_embed. With no embedding head the engine reads the confidence head's probe pool instead, the same pooled residual cells the confidence score is computed from, and returns them as a unit-norm vector of 3072 floats (four probe queries over width 768). It is not a contrastively trained embedding, but it is deterministic, cheap, and a useful similarity signal, so retrieval over a large catalogue is a few lines on any platform:

import numpy as np, needle

agent = needle.Needle(tools=catalogue)
index = np.stack([agent.embed(t["description"]) for t in catalogue])

def shortlist(query, k=5):
    q = np.asarray(agent.embed(query))
    return [catalogue[i] for i in np.argsort(-(index @ q))[:k]]

agent = needle.Needle(tools=shortlist("set an alarm for 6am"))

In C the same call is needle_embed(text, out, capacity); pass a null output to get the dimension. A port that implements the probe pool for the confidence score has this for free.

The practical reason to care: the model was trained with small tool sets, and accuracy falls as the prefix grows past roughly two dozen tools. Large catalogues should be shortlisted before the prompt is built, by this embedding or by any other retriever you already run. When a contrastive embedding head ships in a later checkpoint, the manifest will say so, the engine will shortlist on its own, and --tool-index will cache the catalogue's vectors on disk; nothing about the container changes.

The prompt on the wire

Tokens 0 to 15 are fixed: <pad>, </s>, <s>, <unk>, <|im_start|>, <|im_end|>, <think>, </think>, <tools>, </tools>, <tool_call>, </tool_call>, <tool_result>, </tool_result>, <context>, </context>. The engine builds a prefix once per tool set and caches its KV:

<s><|im_start|>system
date: 2026-09-18 Fri 09:41<|im_end|>
<|im_start|>user
<tools>[{"name":"set_lights","description":"...","parameters":{...}}]</tools>

The system turn is optional and holds session facts (an ISO timestamp becomes the date line above, weekday included); the tools are one minified JSON array. Before rendering, the engine unwraps OpenAI-style {"type":"function","function":{...}} entries, converts tool and property names to snake_case, drops keys the model was not trained on (triggers among them), and keeps an alias table so the calls it returns carry your original names. The first user turn continues the open user message; later turns open their own:

\nturn on the kitchen lights<|im_end|>\n<|im_start|>assistant\n
\n<|im_start|>user\nand the bedroom<|im_end|>\n<|im_start|>assistant\n
\n<|im_start|>tool\n<tool_result>{"ok":true}</tool_result><|im_end|>\n<|im_start|>assistant\n

The model then writes <think> and a one-line derivation, </think>, a newline, and <tool_call>[...]</tool_call>; an empty list is a refusal. The engine forces <think> as the first token of every turn, decodes the call under a byte-level grammar compiled from your schemas, and applies the grounding and confidence gates described in designing tools and confidence. A port can stop at greedy decoding and still get a correct model; the grammar and the gates are what turn a correct model into the contract the guides describe, and they are worth porting second.

The ladder rule

Every depth from 2 to 20 is a trained model, and the engine's --depth N selects blocks without a new file. The rule is endpoint-preserving bisection: start with blocks 0 and 19, then repeatedly add the midpoint of the largest gap between selected blocks, breaking ties toward the left. The order for 20 layers is 0, 19, 9, 14, 4, 6, 11, 16, 2, 7, 12, 17, 1, 3, 5, 8, 10, 13, 15, 18, and depth N keeps the first N in ascending order:

depthblocks
20, 19
40, 9, 14, 19
80, 4, 6, 9, 11, 14, 16, 19
120, 2, 4, 6, 7, 9, 11, 12, 14, 16, 17, 19
160 to 9, 11, 12, 14, 16, 17, 19

An engram site fires only when its layer is active, and the global-attention layers are whichever of 4, 9, 14 and 19 survive. needle build --layers N writes the same selection into a shorter file, with num_layers and the surviving engram sites in the header, so a port that reads the header needs no ladder logic at all to run a sliced archive; it needs the rule only to slice at runtime.

The geometry, for the header

The header's 49 fields are what the engine sizes itself from; the reference config says the same thing in words. For the shipped file: vocabulary 8192 (the output head is the full vocabulary), width 768, 20 blocks, 12 query heads of 48 over 2 key-value heads, value heads of 64, RoPE theta 100000, context 8192, a 1024-token local window with full attention at layers 4, 9, 14 and 19, three causal conv taps on Q, K and V, a Hadamard size of 1024, four residual lanes, engram sites 3, 7, 11, 15 and 19 with orders 2 and 3, 18432 slots and a sub-dimension of 128, an int8 KV cache post-trained at a window of 256, and the shared codebooks of 4, 8 and 16 unit-sphere centroids. The engine runs any geometry the header describes; a port should too, and should refuse the Needle 2 tag 0x05E12A83 with a message rather than a guess.

Community runtimes

If you publish a port, tell us at founders@cactuscompute.com and we will link it here and from the model card, so the community and official paths stay aligned. If a note in this post turns out to be wrong for your implementation, that is a bug in the notes, and we would rather fix it than have you discover it twice.