Views
No views yet
| Parameters | 26M |
| Architecture | Encoder-decoder, pure attention (no FFN) |
| Encoder | 12 layers, GQA (8H/4KV), RoPE, gated residuals |
| Decoder | 8 layers, self-attn + cross-attn, gated residuals |
| d_model | 512 |
| Vocab | 8192 (SentencePiece BPE) |
| Norm | ZCRMSNorm (zero-centered, init=0) |
| Precision | bfloat16 (INT4 QAT during training) |
| Pretraining | 200B tokens on 16x TPU v6e (27hrs) |
| Post-training | 2B tokens of function call data (45mins) |
d=512, 8H/4KV, BPE=8192
┌──────────────┐
│ Tool Call │
└──────┬───────┘
┌┴──────────┐
│ Softmax │
└─────┬─────┘
┌─────┴─────┐
│ Linear (T)│ <- tied
└─────┬─────┘
┌─────┴─────┐
│ ZCRMSNorm │
└─────┬─────┘
┌────────┴────────┐
│ Decoder x 8 │
│┌───────────────┐│
││ ZCRMSNorm ││
││ Masked Self ││
││ Attn + RoPE ││
││ Gated Residual││
│├───────────────┤│
┌──────────────┐ ││ ZCRMSNorm ││
│ Encoder x 12 │─────────────────────>Cross Attn ││
│ │ ││ Gated Residual││
│ ┌──────────┐ │ │└───────────────┘│
│ │ZCRMSNorm │ │ └────────┬────────┘
│ │Self Attn │ │ ┌─────┴─────┐
│ │ GQA+RoPE │ │ │ Embedding │ <- shared
│ │Gated Res │ │ └─────┬─────┘
│ │ │ │ ┌───────┴────────┐
│ │ (no FFN) │ │ │[EOS]<tool_call>│
│ └──────────┘ │ │ + answer │
│ │ └────────────────┘
└──────┬───────┘
│
┌────┴──────┐
│ Embedding │
└────┬──────┘
│
┌────┴──────┐
│ Text │
│ query │
└───────────┘1git clone https://github.com/cactus-compute/needle.git
2cd needle && source ./setup
3needle playground1from needle import load_checkpoint, generate, SimpleAttentionNetwork, get_tokenizer
2
3params, config = load_checkpoint("checkpoints/needle.pkl")
4model = SimpleAttentionNetwork(config)
5tokenizer = get_tokenizer()
6
7result = generate(
8 model, params, tokenizer,
9 query="What's the weather in San Francisco?",
10 tools='[{"name":"get_weather","parameters":{"location":"string"}}]',
11 stream=False,
12)
13print(result)
14# [{"name":"get_weather","arguments":{"location":"San Francisco"}}]1# Web UI (generates data via Gemini, trains, evaluates, bundles result)
2needle playground
3
4# CLI (auto-downloads weights if not local)
5needle finetune data.jsonl@misc{ndubuaku2026needle,
title={Needle},
author={Henry Ndubuaku and Jakub Mroz and Karen Mosoyan and Roman Shemet and Parkirat Sandhu and Satyajit Kumar and Noah Cylich and Justin H. Lee},
year={2026},
url={https://github.com/cactus-compute/needle}
}