❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 丰富的 AI 工具库 -> 每日更新 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦

🚀 "长文本生成进入高铁时代！北通院黑科技让AI写作效率飙升300%"

大家好，我是蚝油菜花。你是否也经历过这些抓狂时刻——

• 👉 写行业白皮书卡在生成环节，5小时只能产出10万Token
• 👉 加速时遭遇"鬼畜循环"，相同段落反复出现
• 👉 换不同模型就要重构代码，开发效率断崖下跌

今天带来的 TokenSwift 正在改写游戏规则！这个由北京通用人工智能研究院打造的文本引擎，凭借多Token并行生成+动态KV缓存两大绝技，实现：

• ✅ 10万Token文本90分钟极速生成（传统方案需5小时）
• ✅ 独创上下文惩罚机制，重复率降低76%
• ✅ 无缝适配1.5B到14B不同规模模型

从万字小说到百页行业报告，现在连科研团队都在用它冲刺论文——你的键盘准备好迎接这场生产力核爆了吗？

🚀 快速阅读

TokenSwift 是一个超长文本生成加速框架，能够在 90 分钟内生成 10 万 Token 的文本。

1. 核心功能：显著缩短超长文本生成时间，支持多种模型架构，生成质量无损。
2. 技术原理：通过多 Token 并行生成与 Token 复用、动态 KV 缓存更新策略、树结构的多候选 Token 验证以及上下文惩罚策略，提升生成效率。

TokenSwift 是什么

TokenSwift 是北京通用人工智能研究院团队推出的超长文本生成加速框架，能够在 90 分钟内生成 10 万 Token 的文本，相比传统自回归模型的近 5 小时，速度提升了 3 倍，生成质量无损。

TokenSwift 通过多 Token 生成与 Token 重用、动态 KV 缓存更新以及上下文惩罚机制等技术，减少模型加载延迟、优化缓存更新时间并确保生成多样性。支持多种不同规模和架构的模型，如 1.5B、7B、8B、14B 的 MHA 和 GQA 模型。

TokenSwift 的主要功能

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

TokenSwift 的技术原理

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

如何运行 TokenSwift

1. 安装

方法 1：使用 pip 安装

pip install tokenswift

方法 2：从源码安装（推荐）

git clone https://github.com/bigai-nlco/TokenSwift.git
cd TokenSwift
conda create -n tokenswift python=3.11
conda activate tokenswift
conda install nvidia::cuda-nvcc
pip install -r requirements.txt
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

2. 模型下载

• TokenSwift-Yarn-Llama-2-7b-128k：https://huggingface.co/TokenSwift/TokenSwift-Yarn-Llama-2-7b-128k
• TokenSwift-Llama-3.1-8B：https://huggingface.co/TokenSwift/TokenSwift-Llama-3.1-8B
• TokenSwift-Qwen2.5-1.5B：https://huggingface.co/TokenSwift/TokenSwift-Qwen2.5-1.5B
• TokenSwift-Qwen2.5-7B：https://huggingface.co/TokenSwift/TokenSwift-Qwen2.5-7B
• TokenSwift-Qwen2.5-14B：https://huggingface.co/TokenSwift/TokenSwift-Qwen2.5-14B
• TokenSwift-DeepSeek-R1-Distill-Qwen-32B：https://huggingface.co/TokenSwift/TokenSwift-DeepSeek-R1-Distill-Qwen-32B

3. 推理

以 LLaMA3.1-8B 为例：

torchrun  --master-port 1111 --nproc_per_node=1 main.py \
    --model_type llama3_1 \
    --ckpt_path your_checkpoint_path \
    --prefill_len 4096 \
    --retrival_max_budget 4096 \
    --gen_len 102400 \
    --gamma 4 \
    --min_p 0.1 \
    --temperature 1.0 \
    --tree_decoding \
    --ngram_topk 20 \
    --penalty 1.2 \
    --penalty_length 1024 \
    --prompt_id 0

  <NOTE: 修改数据和模型路径>

对于其他模型，可以运行 infer_scripts/ 文件夹中的脚本。例如：

bash infer_scripts/r1_qwen_32b.sh

4. 训练指南（可选）

数据集下载

从PG-19训练集中，根据不同的 tokenizer 过滤出大于 8K 的数据。

• PG-19 数据集：https://huggingface.co/datasets/deepmind/pg19

或从以下链接下载处理后的训练数据集：

• llama2-pg19：https://huggingface.co/datasets/TokenSwift/llama2_pg19_train_data
• llama3.1-pg19：https://huggingface.co/datasets/TokenSwift/llama3.1_pg19_train_data
• qwen2.5-pg19：https://huggingface.co/datasets/TokenSwift/qwen2.5_pg19_train_data

如何训练

以 LLaMA3.1-8B 为例：

torchrun --master-port 1111 --nproc_per_node=4 train/train_legacy.py \
    --model_name_or_path /your_model_path/Meta-Llama-3.1-8B \
    --llama_type llama3_1 \
    --data_path /your_data_path/llama3_1_pg19_8k_data \
    --output_dir /your_checkpoint_path/adapter_ckpts_llama3_1 \
    --max_steps 200 \
    --per_device_train_batch_size 3 \
    --gradient_accumulation_steps 10 \
    --save_steps 200 \
    --learning_rate 5e-3 \
    --weight_decay 0.1 \
    --warmup_steps 50 \
    --lr_scheduler_type cosine \
    --logging_steps 5 \
    --report_to tensorboard \
    --bf16 True \
    --medusa_heads 3 \
    --remove-unused-columns false

  <NOTE: 修改数据和模型路径>

对于其他模型，可以运行 train/scripts/ 文件夹中的脚本。例如：

cd train
bash scripts/train_R1_qwen2_5_32b.sh

资源

• 项目主页：https://bigai-nlco.github.io/TokenSwift/
• GitHub 仓库：https://github.com/bigai-nlco/TokenSwift
• HuggingFace 模型库：https://huggingface.co/TokenSwift

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 丰富的 AI 工具库 -> 每日更新 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦