企业官网建设流程全解析

256-token滑动窗口+KV sink：Needle 2工具钉住机制源码分析

【免费下载链接】needle14MB foundation model for tiny devices; phones, wearables, smart home, and robots.项目地址: https://gitcode.com/GitHub_Trending/needle20/needle

Needle 2 是一个面向手机、穿戴设备、智能家居和机器人等极小设备的开源工具调用（Tool Calling）基础模型：整个模型只有一个14MB的二进制文件，跑完整会话仅需约28MB内存。它凭什么在长对话中依然这么省？答案写在 README.md 里的一句话——256-token 滑动窗口 + 把工具定义钉住为 KV Sink（KV 钉位）。这篇文章带你从源码角度拆解 Needle 2 的工具钉住机制是如何实现的。

一张图看懂 Needle 2

模型基于 Simple Attention Network（Hadamard MLP、GQA 注意力、engram 键值记忆、多车道超连接），并用 Cactus Quants 压到 2-bit。下面这张架构图展示了每个 Transformer 块的内部结构，注意其中带缓存的注意力路径——滑动窗口和 KV sink 正是作用在这一层：

为什么小模型最需要「有界内存」

大模型跑长对话，KV cache 会随上下文长度线性增长，内存迟早爆掉。但对 14MB 的小模型来说，预算卡得更死：整个进程只有约 28MB 可用。如果 KV cache 无限增长，模型连多轮工具调用都撑不过去。

Needle 2 的取舍是：只保留最近 256 个 token 的注意力上下文，但永远「看得见」开头被钉住的位置。下图展示了这种极小模型在尺寸-质量曲线上的位置：

256-token 滑动窗口：源码走读

窗口的实现非常克制，核心就在 needle/model/decode.py 的注意力掩码里：

recent = (qpos[:, None] - kpos[None, :]) < cfg.kv_window keep = recent[None] if sink is None else (recent[None] | sink[:, None, :]) causal = causal[None] & keep

三行代码讲清了两件事：

1. 窗口边界：recent只放行「当前位置往前数 256 个 token」以内的 KV；
2. 钉位放行：只要某个位置在sink掩码里为True，它就永远参与注意力，不受窗口限制。

kv_window这个参数从 needle/model/decode.py 的DecodeCfg配置一路传入，Flash Attention 路径也有对应的等价的窗口+sink 掩码（见 needle/model/decode.py）。批量推理路径 needle/model/decode.py 则按 batch 维度为每个样本独立构造 sink 掩码。

KV sink 如何「钉住」工具

所谓KV sink（KV 钉位）：把某些前缀位置的键值对标记为「不可驱逐」，无论上下文滑多远，模型对这些位置始终可见。在 Needle 2 里，被钉住的是工具定义——这样即使对话滚动到几百轮之后，模型依然记得有哪些工具、参数长什么样，随时可以发起工具调用。

源码中 sink 的构造方式：

• 单样本解码时，needle/model/decode.py 先分配一个布尔矩阵sink，再把「文档前缀」段标记为True；
• 训练侧的打包掩码 needle/model/architecture.py 里更直白：sink = prefix > 0，即前缀段内的所有位置都被钉住，掩码为recent | sink；
• 置信度头与对比学习编码头（needle/model/architecture.py）复用同一套window + sink掩码，保证所有头看到的上下文范围一致。

值得注意：Python 参考实现里 _doc_prefix_len 是一个占位钩子（固定返回 0），真正把工具定义段钉成 sink 的动作发生在编译进.cact的引擎里——Python 侧负责定义机制，引擎侧负责执行。

256 这个窗口宽度从哪来：内存预算计算器

窗口不是拍脑袋定的，而是从一个硬性内存预算反推出来的。核心常量在 needle/model/architecture.py：

KV_BUDGET_BYTES = 11 * 1024 * 1024 + 512 * 1024 # 约 11.5MB，留给 KV cache 的总预算 KV_GROUP = 32 # KV 量化按 32 个 token 一组 KV_WINDOW_MIN = 160 # 窗口下限

计算逻辑（needle/model/architecture.py）：

1. 按「每 token 占多少字节」估算：12 层 ×（KV 本体 + 每 32 token 一组的 4 字节量化参数）+ 2 个 engram 站点的额外缓存；
2. 窗口 = 预算 // 每 token 字节数，再向下取整到 32 的倍数（匹配量化组）；
3. 最后取min(预算窗口, 模型声明的 kv_window)，并夹在 160 ~ max_seq_len 之间。

发布版的 Needle 2 检查点声明kv_window=256，预算窗口远大于 256，所以最终生效的就是256。配置字段定义见 needle/model/architecture.py。

这个宽度还会被烧进导出文件：needle build时把effective_kv_window(config)写进.cact文件头（needle/model/finetune.py），头部字段说明在 needle/model/export.py：「kv_window 是模型训练时使用的滑动窗口宽度」。引擎加载后严格按此裁剪内存——这就是「无论对话多长，内存都稳定在 28MB 附近」的工程保证。

快速上手：体验有界内存的 Agent

安装后即可用（权重首次自动下载并缓存）：

pip install cactus-needle

import needle @needle.tool def get_weather(city: str): "Get the current weather for a city." return {"city": city, "temp_c": 27, "sky": "clear"} agent = needle.Needle(tools=[get_weather]) print(agent.run("what's it like in Lagos right now?")["results"])

完整 API（complete()、extract()、置信度门控、工具检索）见 doc/apis.md；用 LoRA 微调并重建自己的.cact（窗口宽度会自动保留）见 doc/finetuning.md。

小结：三个机制，一套有界内存

这套「滑动窗口 + KV sink」的组合拳，本质上是把「模型记住什么」从靠容量变成了靠策略：最近的细节交给窗口，关键的工具契约交给钉位，其余的交给 engram 记忆。对极小设备上的 Agent 来说，这或许是比堆参数更值得借鉴的设计。

【免费下载链接】needle14MB foundation model for tiny devices; phones, wearables, smart home, and robots.项目地址: https://gitcode.com/GitHub_Trending/needle20/needle