TokenSwift

TokenSwift – 超长文本生成加速框架，90分钟内生成10万Token文本

TokenSwift是一款由北京通用人工智能研究院团队开发的超长文本生成加速框架，其强大的功能使得在90分钟内生成10万Token的文本成为可能。与传统自回归模型需要近5小时的生成时间相比，TokenSwift的速度提升达到了三倍，同时不损失生成质量。得益于多Token生成与复用、动态KV缓存更新和上下文惩罚机制等先进技术，TokenSwift有效减少了模型加载延迟，优化了缓存更新时间，确保了生成文本的多样性。

TokenSwift是什么

TokenSwift 是北京通用人工智能研究院团队推出的超长文本生成加速框架，能够在90分钟内生成10万Token的文本。相比传统自回归模型的近5小时，TokenSwift的速度提升达到了三倍，同时保持了生成质量的无损。该框架利用多Token生成与Token重用、动态KV缓存更新以及上下文惩罚机制等技术，显著降低了模型加载延迟，优化了缓存更新时间，并确保生成文本的多样性。它支持多种规模和架构的模型，包括1.5B、7B、8B、14B的多头注意力（MHA）和分组查询注意力（GQA）模型。

TokenSwift的主要功能

• 加速超长文本生成：TokenSwift可以显著缩短超长文本的生成时间，传统自回归模型生成10万Token的文本可能需要近5个小时，而TokenSwift仅需90分钟，效率提高了300%。
• 保持生成质量：在加速生成的同时，TokenSwift确保了文本输出的质量和多样性，原始模型的输出质量得以保留。
• 兼容多种模型：TokenSwift支持不同规模和架构的多种模型，包括1.5B、7B、8B和14B的模型，以及多头注意力（MHA）和分组查询注意力（GQA）架构。

TokenSwift的技术原理

• 多Token并行生成与复用：TokenSwift借鉴了Medusa等先进方法，引入额外的线性层，使得模型能够在一次前向传播中同时生成多个草稿Token。系统基于生成文本中的n-gram频率信息，自动检索并复用高频短语，从而减少模型的重新加载次数，提高整体效率。
• 动态KV缓存管理策略：TokenSwift在KV缓存管理方面采用动态更新策略，生成过程中保留初始KV缓存，并根据Token的重要性对后续缓存进行有序替换。
• 树结构的多候选Token验证：为确保生成结果与目标模型预测的一致性，TokenSwift引入了树形注意力机制。通过构建包含多个候选Token组合的树形结构，采用并行验证的方式，从中随机选择最长且有效的n-gram作为最终输出，确保生成过程的无损性和多样性。
• 上下文惩罚机制：为了进一步减少重复生成现象，TokenSwift设计了一种上下文惩罚方法。在生成过程中对近期生成的Token施加惩罚，使得模型在选择下一Token时更倾向于多样化输出，有效降低重复现象的发生。

TokenSwift的项目地址

• 项目官网：https://bigai-nlco.github.io/TokenSwift/
• Github仓库：https://github.com/bigai-nlco/TokenSwift
• HuggingFace模型库：https://huggingface.co/TokenSwift
• arXiv技术论文：https://arxiv.org/pdf/2502.18890

TokenSwift的应用场景

• 内容创作与文案生成：TokenSwift在短时间内生成高质量文本内容，非常适合内容创作者、广告公司和媒体机构，显著提高工作效率。
• 智能客服与机器人：在智能客服和机器人应用中，TokenSwift能迅速生成详尽的回答，尤其在处理复杂问题或提供长篇解释时，能带来更高效、更连贯的对话体验。
• 学术研究与论文撰写：对于学术研究和论文撰写，TokenSwift能帮助研究人员快速生成文献综述、实验报告或论文草稿，节省大量时间和精力。
• 代码生成与编程辅助：在编程领域，TokenSwift可用于生成复杂的代码片段或文档，帮助开发者迅速构建和优化代码。