十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

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

256-token滑动窗口+KV sink:Needle 2工具钉住机制源码分析 256-token滑动窗口KV sinkNeedle 2工具钉住机制源码分析【免费下载链接】needle14MB foundation model for tiny devices; phones, wearables, smart home, and robots.项目地址: https://gitcode.com/GitHub_Trending/needle20/needleNeedle 2 是一个面向手机、穿戴设备、智能家居和机器人等极小设备的开源工具调用Tool Calling基础模型整个模型只有一个14MB的二进制文件跑完整会话仅需约28MB内存。它凭什么在长对话中依然这么省答案写在 README.md 里的一句话——256-token 滑动窗口 把工具定义钉住为 KV SinkKV 钉位。这篇文章带你从源码角度拆解 Needle 2 的工具钉住机制是如何实现的。一张图看懂 Needle 2模型基于 Simple Attention NetworkHadamard 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三行代码讲清了两件事窗口边界recent只放行「当前位置往前数 256 个 token」以内的 KV钉位放行只要某个位置在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 sinkKV 钉位把某些前缀位置的键值对标记为「不可驱逐」无论上下文滑多远模型对这些位置始终可见。在 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.pyKV_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按「每 token 占多少字节」估算12 层 ×KV 本体 每 32 token 一组的 4 字节量化参数 2 个 engram 站点的额外缓存窗口 预算 // 每 token 字节数再向下取整到 32 的倍数匹配量化组最后取min(预算窗口, 模型声明的 kv_window)并夹在 160 ~ max_seq_len 之间。发布版的 Needle 2 检查点声明kv_window256预算窗口远大于 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-needleimport 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(whats it like in Lagos right now?)[results])完整 APIcomplete()、extract()、置信度门控、工具检索见 doc/apis.md用 LoRA 微调并重建自己的.cact窗口宽度会自动保留见 doc/finetuning.md。小结三个机制一套有界内存机制所在位置作用256-token 滑动窗口needle/model/decode.py注意力只看最近 256 个 token上下文成本恒定KV sink 钉位needle/model/architecture.py工具定义段永不被窗口驱逐长对话中工具始终可用内存预算反推窗口needle/model/architecture.py11.5MB KV 预算 32-token 量化组把窗口宽度与总内存锁死这套「滑动窗口 KV sink」的组合拳本质上是把「模型记住什么」从靠容量变成了靠策略最近的细节交给窗口关键的工具契约交给钉位其余的交给 engram 记忆。对极小设备上的 Agent 来说这或许是比堆参数更值得借鉴的设计。【免费下载链接】needle14MB foundation model for tiny devices; phones, wearables, smart home, and robots.项目地址: https://gitcode.com/GitHub_Trending/needle20/needle创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表