Needle 3 for Sentis (com.sky.sentis.needle3)

Unity Sentis (com.unity.ai.inference 2.6.1) port of Cactus-Compute/needle3, a 121M-parameter on-device tool-calling model. It runs fp32 and unquantised, from the needle3.safetensors master weights (fp16 → fp32). It does not use the 2-bit needle3.cact that the vendor's C++ engine ships.

using Needle3;
using Unity.InferenceEngine;

using var needle = new Needle3Model(backend: BackendType.GPUCompute);  // loads StreamingAssets/Needle3
needle.Warmup();
var r = needle.Complete(toolsJson, "dim the living room to 30");
// r.Calls[0] -> set_light({"room":"living room","brightness":30}),  r.Confidence ~ 0.96
// r.Reasoning, r.ToolCallJson ("[]" = refusal), r.GeneratedIds, r.PrefillMs / r.DecodeMs

Adding tools

Tools are a JSON array of schemas passed as toolsJson. There is no registration step: to add a tool, add an entry to the array. Each entry has a name, a description and a JSON Schema parameters object. This is the shape the cactus-needle Python package produces with @needle.tool.

const string Tools = @"[
  {""name"": ""set_light"", ""description"": ""Turn a room's lights on or off, or dim them."",
   ""parameters"": {""type"": ""object"", ""properties"": {
       ""room"": {""type"": ""string"", ""description"": ""Room name, e.g. living room.""},
       ""brightness"": {""type"": ""integer"", ""description"": ""0-100 percent.""}},
     ""required"": [""room"", ""brightness""]}},
  {""name"": ""log_expense"", ""description"": ""Record an expense."",
   ""parameters"": {""type"": ""object"", ""properties"": {
       ""amount"": {""type"": ""number"", ""description"": ""Amount in dollars.""},
       ""category"": {""type"": ""string"", ""enum"": [""food"", ""transport"", ""rent""]},
       ""note"": {""type"": ""string""}},
     ""required"": [""amount"", ""category""]}}
]";

var r = needle.Complete(Tools, "spent 12 dollars on lunch");
// r.Calls -> log_expense({"amount":12,"category":"food","note":"lunch"})
foreach (var call in r.Calls)
{
    switch (call.Name)
    {
        case "log_expense":
            var amount = ((JsonNumber)call.Arguments["amount"]).ToDouble();  // numbers arrive as JsonNumber
            var category = (string)call.Arguments["category"];
            call.Arguments.TryGetValue("note", out var note);                // optional: may be absent
            break;
        // case "set_light": ...
    }
}

The tools are minified with json.dumps(separators=(",", ":")) rules and rendered into the prompt. Whitespace and indentation in your string don't matter. OpenAI-style {"type": "function", "function": {...}} entries are unwrapped, and triggers keys are dropped.

Designing tools (from the vendor's guide, Designing tools for Needle):

• One tool per action, named in snake_case with words users would say. Polar pairs such as lock_door / unlock_door are separate tools, or one enum. Don't make the model infer polarity from the description.
• Descriptions state what the tool does, not instructions to the model. Put synonyms and value formats there, for example "ISO date" or "City, ST".
• Use enum for closed sets, with values users actually say. ["increase", "decrease"] works; ["inc", "dec"] doesn't.
• List only the required parameters in required. Anything the request may not mention should be optional.
• Keep the set small. Every tool goes into the prompt, so prefill time grows with the catalogue, and accuracy drops past about two dozen tools. This port does not shortlist tools by embedding retrieval the way the vendor engine does. For a large catalogue, pick the relevant few yourself per request, or use two turns: one to choose the tool, then one with only that tool.
• Put environment facts in system, not in tool descriptions. For example: needle.Complete(Tools, query, system: "date: 2026-10-17 Sat 09:41").
• No call means no tool fits. An empty r.Calls (with r.ToolCallJson == "[]") is the model declining. Treat it as a refusal; don't fall back to a guess.

⚠ The vendor engine compiles schemas into a decoding grammar, so its bounds (minimum, maxLength, …) and enum constraints are enforced while decoding, and triggers force a call. This port decodes greedily without that grammar. Validate the arguments in your handler (range, enum membership, presence) before acting on them.

Layout

Path	What
Model~/needle3_step.sentis	fp32 incremental step graph, serialized Sentis model (512 MB)
Model~/needle3_confidence.sentis	fp32 confidence head (probe pool over the 21 residual depths)
Model~/tokenizer.json	SentencePiece BPE pieces/scores/types for the C# tokenizer
Model~/fixtures/	golden data from the PyTorch port (== JAX reference) used by the tests
Runtime/	Needle3Model (runner), Needle3Tokenizer, Needle3Prompt (prompt, engram hash, parser), MiniJson
Editor/Needle3Baker.cs	copies Model~ into Assets/StreamingAssets/Needle3 (menu Tools ▸ Needle3 ▸ Install Models to StreamingAssets)
Tests/Editor/	host parity + model parity/end-to-end tests (CPU and GPUCompute)
Tools~/	Python export and parity scripts (they produce the ONNX the .sentis files were serialized from)

Model~ is hidden from the asset pipeline and is not included in player builds. Run the install menu once, or call Needle3Baker.Bake(), so the runtime can load the .sentis files from StreamingAssets.

Graph contract (needle3_step)

S = new tokens this call, P = cached tokens, T = P + S. The C# loop owns all the state.

Input	Shape	Notes
input_ids	[1,S] int
position_ids	[S] int	absolute positions P..T-1 (RoPE θ=1e5, gathered from an in-graph table, max 8192)
engram_idx	[1,S,6] int	table*18432 + hash (uint32 n-gram hash computed host-side, Needle3Prompt.EngramRows)
ngram_ok	[1,S,6] float	1 when position ≥ order-1
mask_global / mask_local	[1,1,S,T] float	additive 0/-1e9. Local = 1024 sliding window; global layers are 4, 9, 14, 19
past_key_i / past_value_i	[1,2,P,48] / [1,2,P,64]	post-RoPE GQA cache, 20 layers
past_conv_i	[1,2,800]	last 2 pre-tap q|k|v rows (3-tap causal conv)
past_eng_s	[1,9,768]	last 9 pre-tap engram value rows (4 taps, dilation 3), 5 sites

Outputs: logits [1,S,8192], cells [1,S,21,768] (lane-mean residual per depth, input to the confidence head) and present_*, which mirror every past_*.

Prompt

This follows render_example in needle/model/finetune.py:

<s>[<|im_start|>system\n{system}<|im_end|>\n]<|im_start|>user\n<tools>{minified tools JSON}</tools>\n{query}<|im_end|>\n<|im_start|>assistant\n<think>

<think> is forced, as in the vendor engine. Decoding is greedy and stops at <|im_end|> or </s>. OpenAI-style {"type":"function","function":{…}} entries are unwrapped before minifying.

Verification (Sentis 2.6.1, Unity 6000.3.23f1, Apple Silicon)

Edge	Result
JAX (original) vs PyTorch port	logits max|Δ| 2.4e-4, cos 1.0; full and incremental KV paths; confidence logit identical
PyTorch vs ONNX Runtime	prefill logits max|Δ| 9.2e-5; greedy tokens identical; confidence 3.41743 vs 3.41743
ORT/torch fixture vs Sentis	CPU logits max|Δ| 9.2e-5, GPUCompute 7.6e-5; cells rel ≤ 1.8e-7
End to end (6 cases)	greedy ids token-identical to the reference on CPU and GPUCompute: single call, two calls in order, refusal [], system prompt

Measured in the editor test run (warm, M-series Mac):

Backend	Prefill (~150 tok)	Decode
GPUCompute	~100 ms	~21 tok/s
CPU (Burst)	~11 s	~8 tok/s

Use GPUCompute where compute shaders are available. CPU prefill is slow because the graph has ~8.7k layers, many of them tiny (20-iteration Sinkhorn per layer, Monarch factors).

Differences from the vendor engine

• fp32, no quantisation. The vendor engine runs CQ2/CQ4 weights with int8 activations and an int8 KV cache.
• No grammar-constrained decoding. The engine's byte-level schema grammar is closed source. Here the output is plain greedy text, and the JSON is parsed afterwards (Calls is empty if it fails to parse).
• No auto date facts, snake-casing or tool retrieval. Pass any session facts as system.
• Confidence = sigmoid of the shipped confidence head over the whole turn (prompt + generated tokens). The engine's act/confirm/refuse gates are not reproduced.
• KV window: full attention up to 8192 tokens (1024 sliding window on local layers), as in the Python reference. The engine's 256-token KV window is not used.

Weights, tokenizer and engine source are © Cactus Compute, Apache-2.0 (see LICENSE.md).