KTransformers 长上下文稀疏注意力实战单张 24GB GPU 跑通 1M Token KVCache 的完整方案【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers本文以 KTransformers 官方长上下文技术文档为主体系统讲解“KVCache 超长上下文”这一场景的问题本质、稀疏注意力的数学原理、H2O/InfLLM/Quest/SnapKV 等代表性工作的思路对比以及 KTransformers 自研的 CPU 稀疏注意力框架的架构、全部可配置参数与实验数据并结合仓库源码给出long_context模式的配置与运行方法。读完本文你将理解为什么 1M token 上下文在 24GB 显存的消费级 GPU 上可以实跑并掌握如何复现 128K/1M 场景下的加速配置。1. 问题背景超长上下文的 KVCache 内存墙训练更大模型与支持更长文本序列是当前公认的通向 AGI 的两个方向。KTransformers 在降低万亿参数 MoE 模型本地推理门槛之后长上下文推理是其主打的第二类展示场景。ChatGLM 与 InternLM 均发布了支持 100 万 token 上下文的开源模型本文以 InternLM2.5-7B-Chat-1M 为例介绍如何利用注意力的稀疏性在 CPU/GPU 异构系统上加速长文本推理。核心矛盾在于模型权重本身并不大但超长上下文的 KVCache 足以压垮本地硬件。以 InternLM2.5-7B-Chat-1M 为例模型权重仅需 15.49GB 显存而完整 1M token 的 KVCache 需要额外 145.49GB 内存明显超出本地用户能力即使借助 llama.cpp 的 KVCache Offload 把 KVCache 卸载到 CPU/DRAM 勉强跑通性能也不可用——因为每生成一个 token 都必须完整扫描一遍整个 KVCache。2. 数学原理Softmax 带来的注意力稀疏性幸运的是多项研究观察到推理阶段注意力分布呈稀疏形态。SparQ 基于 LLaMA 7B 的实验统计显示3k 上下文中注意力得分较高的 token 不足 1%H2O、Quest、InfLLM、SnapKV 等工作得出了类似结论。KTransformers 团队也在 InternLM 2.5-7B-1M 上通过长文本实验进一步验证虽然比例没到 1% 那么极端但由于 attention 中 softmax 天然存在“头部聚焦”效应只要能提前识别哪些 token 的注意力得分高扫描不到 5% 的 token 就基本能复刻原始结果。于是问题收敛为如何在不全量扫描的前提下快速定位这些高注意力 token。这正是 KTransformers CPU 稀疏注意力框架要解决的核心问题。3. 相关工作Prune 还是 Retrieval团队系统调研了近年的 KVCache 稀疏选择工作可以归为两大流派H2O 及后续工作Prune 流派最早指出推理注意力稀疏、只需保留 5% KVCache 的结论后续工作在此基础上设计更精细的 token 选择策略。这类方法对单条查询推理合理但如 Mooncake 项目的探索所指出的未来趋势是尽可能预计算可复用的 KVCache再用它回答不同问题“一次计算、多次复用”。因此 KTransformers 倾向于不删除 KVCache 中的 token或至少不删掉大部分保证不同问题可以各自关注上下文的不同部分。InfLLMRetrieval 流派其框架与上述理念高度契合被 KTransformers 选作核心参考。InfLLM 的核心设计将过长的上下文 KVCache 存入外部记忆模块Memory Units每步计算时只检索最相关的语义信息参与运算缓解“过长上下文导致注意力被无关噪声稀释”的问题外部记忆以语义块相邻 token 组成组织计算时采用滑动窗口机制——只选取上下文头部的块Initial Tokens、靠近当前 token 的块Local Tokens以及与当前 token 语义相似度最高的少量块参与注意力计算为高效检索InfLLM 在每个块内选取注意力得分 $$r_m$$ 最高的若干代表 token用相似度方程计算当前 token 与每个语义块的相关性见 InfLLM 公式图。与 H2O 相比InfLLM 的两个关键区别KVCache 不丢弃而是存放在内存中并在推理时动态加载到 GPUKVCache 以块为粒度管理每块选少量代表 token。其外部记忆单元可卸载到 CPU/DRAM 甚至 SSD按具体问题选取不同部分参与计算。3.1 Quest块级代表 token 与上界估计Quest 同样以块为粒度管理 token。它分析了 H2O 与全注意力下关键 token 的召回率发现 H2O 对 Top-10 注意力 token 的召回率仅约 50%关键信息丢失过多。为此 Quest 从每个块中选两个“代表 token”做检索Prefill 阶段每个 KVCache 块记录每个通道的最大值与最小值即图中标注的 Reduced Keys含逐元素 min key 与 max key注意力计算阶段当前 query 向量分别与每个块的 max key 和 min key 做点积对每个通道取两个乘积向量的最大值再求和作为该块相关性的上界根据相关性得分选出 top-k 块参与注意力计算。Quest 没有考虑异构架构假设全部 KVCache 仍可放进内存仅用稀疏注意力加速推理。最终 Quest 实现了 7.03x 注意力计算加速与 2.23x 端到端推理时延改善。3.2 SnapKV 及其他工作SnapKV 在 prefill 阶段一次性保留两部分 token图中橙色与绿色区段与 InfLLM 的区别只在于中间 token 的选择方式SnapKV 在token 级而非块级做选择打分方式与 H2O 类似即 $$softmax(\frac{qk^T}{\sqrt{d_k}})$$但按列求和时只统计最后绿色窗口对应 InfLLM 的 Local Tokens内的行并额外引入 pooling 操作论文解释是为了让召回 token 保留更完整的语义信息。SnapKV 是一次性选择选择完成后只用被选 token 做注意力其余 KVCache 直接丢弃。其他相关工作PyramidKV观察到注意力分数在层间呈金字塔分布——低层注意力分布广泛高层少数关键 token 的得分越来越突出。因此 PyramidKV 给低层分配更多 KVCache 容量、高层更少。MagicPiG基于 Locality-Sensitive HashingLSH的动态 KVCache 管理——先用 SnapKV 选一部分重要 token 存 GPU其余 KVCache 放内存利用 LSH 在高维空间检索中的高效性与 CPU 多线程能力检索与当前 query 相似的内存 KVCache 并载入。相比 InfLLM/Quest/SnapKVMagicPiG 不需要扫描全部代表 token 再选 top-k而是借助 LSH 的数学性质以低开销、高速度模拟注意力得分并定位重要 KVCache。4. KTransformers CPU 稀疏注意力框架4.1 设计要点从上述工作中KTransformers 提炼出四条关键结论注意力权重分布稀疏无用的 KVCache 可能引入噪声反而降低推理性能推理阶段 KVCache 驱逐策略的通行做法是保留 prompt 头尾的 token为中间部分设计选择算法模型性能的主要影响因素是能否准确识别关键 token中间部分按块管理可提升内存换入换出与注意力计算效率且小块粒度并不比 token 级表现更差各注意力层在推理时关注的 token 不同不同层分配到的 KVCache 容量也应不同。在此之上KTransformers 实现了通用的推理期稀疏 CPU 注意力算子框架。整体流程为Prefill 阶段采用 chunked prefill每次只把一层 KVCache 载入 GPU 计算算完存回 CPU/DRAMDecode 阶段不再来回换入换出 KVCache稀疏注意力算子直接运行在 CPU 上。这显著降低了最小 GPU 显存需求使本地 128K 甚至 1M token 上下文成为可能。生成阶段的完整框架如下4.2 KVCache 切分Initial / Context / LocalKVCache 按块组织完整输入 prompt 被切分为三个可配置部分Initial序列头部全程全量参与注意力Context中间部分稀疏检索Local序列尾部全程全量参与注意力。头尾全量保留的设计依据来自 streamingLLM、Minference 等论文发现的“注意力汇聚attention sinks”现象注意力权重往往显著偏向序列开头与结尾。4.3 Context 块划分与代表 token 选择中间 Context 部分沿用 InfLLM 思路按可配置固定 token 数划分为块每块可选 1k 个 token 作为代表 token推理时基于这些代表 token 挑选需要参与注意力的 Context 块。框架实现了五种代表 token 选择方法方法说明Max块内多个 token 逐通道取最大值拼成该块的 1 个代表 tokenMean块内多个 token 逐通道取均值拼成代表 tokenQuest取块内逐通道最大值与最小值作为代表 token代表 token 数固定为 2Dynamic按特定方法计算每个 token 的累计注意力得分每块选 top-k 最高分 token 作为代表类似 InfLLM 但做了简化Fix在块内按固定间隔选取 token代表 token 确定后使用 InfLLM 的相似度方程计算输入 X 与每个块 B 的 k 个代表 token 的相似度只选出 top $$r_k$$ 个块参与注意力计算其中 $$l_P$$ 为历史 token 长度。关于 InfLLM 原方法文档说明了两点工程取舍它需要在 prefill 阶段为每个 token 计算代表性得分、再依此选块代表 token这一操作对 prefill 实现是侵入式修改难以与其他方法整合且实测中多数场景下用其他方法组合即可达到类似甚至更好的效果。因此该方式最终未并入框架但Dynamic方法保留了其简化版。4.4 预选择Preselect借鉴 SnapKV 的可选模块框架初版对每个生成 token 都重新选择最相关的 KVCache 块。这既有明显的检索开销又发现频繁更换所选块并不会带来更好结果。例如 kvretrieval 数据集上模型回答常常前半正确、后半错误——由于答案是无意义的长字符串这说明前期 token 生成时选对了 KVCache 块后期却选错了。为此框架集成了 SnapKV 的方法推理开始前基于 question 的上下文 token 注意力得分预先选择一批相关 KVCache 块后续推理中块的选择被限制在这批预选范围内。在 kvretrieval 数据集上该机制可以 100% 选中包含正确答案的块。需要强调该方法严格依赖 Benchmark Prompt 的结构在复杂文档理解与生成等其他场景不保证最优因此被设计为可选模块。最终框架与可配置参数如下4.5 完整可配置参数一览threads_numCPU 线程数block_sizeKVCache 块大小local_windows_lenprompt 尾部窗口大小preselect_block_count预选块数量second_block_count预选后每次再选出的块数preselect_block是否启用预选择token_stepKVCache 块重选的 token 间隔layer_stepKVCache 块重选的层间隔dense_layer_num前若干层不做 KVCache 选择、全量参与head_select_modeSEPARATEGQA 场景下每个 kv_head 独立选择/SHARED所有 kv_head 联合选择representative_type配置文件名anchor_type代表 token 选择方法representative_num配置文件名anchor_num代表 token 数量通过调整这些参数可以复现多种 KVCache 驱逐/压缩方法例如block_size设为 1 且preselect_block为 True即得到无 pooling 的 SnapKV 版本representative_type设为Quest、preselect_block设为 False、head_select_mode设为SEPARATE即复现 Quest 方法。4.6 框架伪代码def preselect_block(local_q, kvcache): key_states kvcache.keycache attn_scores torch.matmul( local_q, key_states.transpose(2, 3) ) / math.sqrt(head_dim) attn_scores attn_mask attn_scores nn.functional.softmax( attn_scores, dim-1, dtypetorch.float32 ).to(query_states.dtype) vote attn_scores[..., initial_size:-local_size:, :].sum(dim-2) pool_vote pool1d(vote, kernel_sizekernel_size, paddingkernel_size//2, stride1) indices pool_vote.topk(max_capacity_prompt - local_size, dim-1).indices kv_cache_block_indices find_representative_tokens_block(indices) kvcache_after_preselected kvcache[kv_cache_block_indices] ... return kvcache_after_preselected def get_representative_tokens(): # 根据 representative_type 计算每个块的代表 token return ... def decode_attention(query, key, value): # 每 token_steps 个 token 选择一次 token_steps 4 # 每 layer_steps 层选择一次 layer_steps 4 for token_idx in range(max_new_tokens): for layer_idx in range(config.num_hidden_layers): if token_idx % token_steps ! 0 or layer_idx % layer_steps ! 0: # 本轮当前层无需重选时复用 kvcache 中保存的历史选择结果 kvcache_after_retrieval history_kvcache_after_retrieval[layer_idx//layer_steps] else: # 否则用当前轮当前层的 query 重新选择 kvcache kvcache_after_retrieval retrieval_kvcache(query, kvcache) # 保存进 kvcache 历史选择结果 history_kvcache_after_retrieval[layer_idx//layer_steps] kvcache_after_retrieval # 计算注意力 output attn(query, kvcache_after_retrieval) yield output # 模型 prefill若需预选择local_q 仍需保存 local_q, KVCache model.prefill(input_ids) if preselect_block: # 预选择轮 KVCache preselect_block(local_q, kvcache) # 找到每个块的代表 token block_representative_tokens get_representative_tokens( kvcache, config.representative_type ) # model generate decode_attention(query, key, value)伪代码体现了两个关键的工程优化token_steps/layer_steps让“选择”与“注意力计算”解耦——并非每个 token 的每层都重新检索而是按间隔重选、其余时间复用历史选择结果history_kvcache_after_retrieval按层组缓存直接降低了检索开销。5. 实验结果5.1 基线配置实验初始采用如下基础配置后续在扩展框架上进一步优化max_seq_len: 256000 # KVCache 长度 block_size: 128 # KVCache 块大小 local_windows_len: 4096 # 长度 local_windows_len 的 KVCache 存放在 GPU 上 second_block_count: 96 # 预选后每次选择的 KVCache 块数若 preselect_block_count则直接使用预选块 threads_num: 64 # CPU 线程数 representative_type: DYNAMIC # KVCache 块代表 token 选择方法 kv_type: FP16 dense_layer_num: 0 # 前若干层不填充/不选择 KVCache representative_num: 1 # 每个 KVCache 块内代表 token 数量 preselect_block: False # 是否启用预选择 head_select_mode: SHARED # 所有 kv_head 联合选择 preselect_block_count: 0 # 预选块数量 layer_step: 1 # 每几层选择一次 token_step: 1 # 每几个 token 选择一次在 opencompass 框架下passkey 数据集是在冗余文本的不同深度插入一小段数字kvretrieval 则是在随机生成的键值对中查找匹配项。128K 场景下原始模型与加速后 KTransformers 的对比Single needle retrieval zh 128k / passkey / kvretrieval方案Single needle retrieval zh 128kpasskeykvretrieval原始模型99.8910021.0KTransformers每个生成 token 重选 KVCache 块10010015.40简单数据集上两者均接近满分生成速度从 llama.cpp 的 4.86 tokens/s 提升到 KTransformers 的 27.49 tokens/s最高 5.65x 加速。较难的 kvretrieval 上初始配置明显掉分需要下一节的优化策略来弥补甚至超越原始模型。另外团队也用该配置框架复现了 Quest 的结果。由于 InternLM2.5-7B-Chat-1M 使用 GQAGrouped Query Attention而 Quest 论文主要针对 MHA 模型优化实测结果不理想Quest 官方也承认需要进一步支持 GQA故不再展开。5.2 预选择 稠密层超越原始模型前文提到逐 token 重选会导致文本越长越容易偏离、后期选到不同块。引入 SnapKV 式预选择后kvretrieval 分数从 15.4 提升到 24.2超过原始模型全注意力的 21 分。进一步发现 LLM 前几层 KVCache 的稀疏化效果不显著因此将前两层设为全量复用 KVCache最终取得24.4分相对原始 21.2 分。1M 数据集上的 needle-in-a-haystack 测试中KTransformers CPU 稀疏注意力框架不仅复现了原始模型报告分数还将准确率从 89.31 提升到92.88且推理速度达到 llama.cpp 的近 10 倍。5.3 与 llama.cpp 的系统级对比在 4090D 服务器上llama.cpp 配置为 KVCache 存 CPU/DRAM、全部计算在 GPU以 Single Needle Retrieval 数据集对比在保持100% 答题正确率的前提下prefill 提速 20.694.1 倍推理提速1.27.1 倍。prefill 差距大的主因llama.cpp 开启 KVCache offload 后在 CPU 上做注意力计算而长文本场景注意力不仅计算重还占据绝大部分计算时间KTransformers 则通过模板注入框架实现 GPU 逐层 Chunk Prefill。推理速度方面KTransformers 随 prompt 增长保持平稳近似水平线而 llama.cpp 随 prompt 变长逐渐变慢——因为 KTransformers 只选取最重要的 16K KVCache 块参与推理速度始终相当于 llama.cpp 处理 16K prompt 的水平不随上下文增长退化至少在这些测试数据集上如此。6. 使用指南适用范围长上下文功能目前仅通过local_chat.py接口支持仓库中对应实现位于 archive/ktransformers/local_chat.pyserver 接口的集成当时仍在开发中且long_context模式断言要求模型架构为LlamaForCausalLMInternLM2.5-1M 正是以 Llama 架构转换发布的。从 源码 可以看到约束与精度设置if mode long_context: assert config.architectures[0] LlamaForCausalLM, only LlamaForCausalLM support long_context mode torch.set_default_dtype(torch.float16)配套的优化规则文件为 Internlm2_5-7b-Chat-1m.yamlembed_tokens的 generate/prefill 均放在 CPULlamaModel与self_attn放在 CUDA 上执行——这正是“decode 在 GPU 做轻量计算、KVCache 主体驻留 CPU/DRAM”的算子级落点。团队已将对应的模型 config、GGUF 与 tokenizer 上传到公开的模型仓库按 InternLM2.5 转 Llama 1M 命名的仓库搜索即可获取。6.1 配置文件首次运行local_chat.py后会在~/.ktransformers下自动创建config.yaml由 Config 单例 从 仓库模板 拷贝而来。长上下文相关配置chunk_size: 4096 # prefill chunk size max_seq_len: 100000 # KVCache 长度 block_size: 128 # KVCache 块大小 local_windows_len: 4096 # 长度 local_windows_len 的 KVCache 存放在 GPU 上 second_select_num: 96 # 预选后每次选择的 KVCache 块数若 preselect_block_count则直接使用预选块 threads_num: 64 # CPU 线程数 anchor_type: DYNAMIC # KVCache 块代表 token 选择方法 kv_type: FP16 dense_layer_num: 0 # 前若干层不填充/不选择 KVCache anchor_num: 1 # 每个 KVCache 块内代表 token 数量 preselect_block: False # 是否启用预选择 head_select_mode: SHARED # 所有 kv_head 联合选择 preselect_block_count: 96 # 预选块数量 layer_step: 1 # 每几层选择一次 token_step: 1 # 每几个 token 选择一次注意源码中配置项命名为anchor_type/anchor_num/second_select_num见 config.py与文档概念名representative_type/representative_num一一对应。max_seq_len必须在~/.ktransformers/config.yaml中调整运行时 local_chat.py 会断言max_seq_len 输入长度 max_new_tokens。6.2 不同上下文长度所需内存上下文长度4K32K64K128K512K1MDRAM 占用 (GB)0.54.298.5817.168.7145.49请根据自身 DRAM 容量选择合适的max_seq_len。示例命令python local_chat.py --model_path/data/model/internlm2_5_to_llama_1m --gguf_path/data/model/internlm2_5_to_llama_1m --max_new_tokens500 --cpu_infer10 --use_cuda_graphTrue --modelong_context --prompt_file/path/to/file若已通过prompt_file指定了输入文本终端显示chat:后直接按回车即可开始推理。7. 小结KTransformers 的长上下文方案可以概括为三句话Prefill 阶段 GPU 逐层 chunk 计算、KVCache 落盘 CPU/DRAMDecode 阶段稀疏注意力直接跑在 CPU 上块级代表 token 间隔重选 可选的 SnapKV 预选择构成完整的选择策略。这一设计让单张 24GB GPU 可以在 128K 场景下以 7.1 倍于 llama.cpp 的速度、1M 场景下以约 16 tokens/sllama.cpp 近 10 倍的速度维持原生精度推理并且在 kvretrieval、needle-in-a-haystack 等基准上达到甚至超过原始模型全注意力的分数。其价值不仅在于具体的加速数字更在于提供了一个参数化、可复现 SnapKV/Quest 等多种方法的统一稀疏注意力框架为后续结合 MInference 等高性能稀疏 prefill 方法留出了扩展空间。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 KTransformers 长上下文稀疏注意力实战：单张 24GB GPU 跑通 1M Token KVCache 的完整方案
KTransformers 长上下文稀疏注意力实战单张 24GB GPU 跑通 1M Token KVCache 的完整方案【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers本文以 KTransformers 官方长上下文技术文档为主体系统讲解“KVCache 超长上下文”这一场景的问题本质、稀疏注意力的数学原理、H2O/InfLLM/Quest/SnapKV 等代表性工作的思路对比以及 KTransformers 自研的 CPU 稀疏注意力框架的架构、全部可配置参数与实验数据并结合仓库源码给出long_context模式的配置与运行方法。读完本文你将理解为什么 1M token 上下文在 24GB 显存的消费级 GPU 上可以实跑并掌握如何复现 128K/1M 场景下的加速配置。1. 问题背景超长上下文的 KVCache 内存墙训练更大模型与支持更长文本序列是当前公认的通向 AGI 的两个方向。KTransformers 在降低万亿参数 MoE 模型本地推理门槛之后长上下文推理是其主打的第二类展示场景。ChatGLM 与 InternLM 均发布了支持 100 万 token 上下文的开源模型本文以 InternLM2.5-7B-Chat-1M 为例介绍如何利用注意力的稀疏性在 CPU/GPU 异构系统上加速长文本推理。核心矛盾在于模型权重本身并不大但超长上下文的 KVCache 足以压垮本地硬件。以 InternLM2.5-7B-Chat-1M 为例模型权重仅需 15.49GB 显存而完整 1M token 的 KVCache 需要额外 145.49GB 内存明显超出本地用户能力即使借助 llama.cpp 的 KVCache Offload 把 KVCache 卸载到 CPU/DRAM 勉强跑通性能也不可用——因为每生成一个 token 都必须完整扫描一遍整个 KVCache。2. 数学原理Softmax 带来的注意力稀疏性幸运的是多项研究观察到推理阶段注意力分布呈稀疏形态。SparQ 基于 LLaMA 7B 的实验统计显示3k 上下文中注意力得分较高的 token 不足 1%H2O、Quest、InfLLM、SnapKV 等工作得出了类似结论。KTransformers 团队也在 InternLM 2.5-7B-1M 上通过长文本实验进一步验证虽然比例没到 1% 那么极端但由于 attention 中 softmax 天然存在“头部聚焦”效应只要能提前识别哪些 token 的注意力得分高扫描不到 5% 的 token 就基本能复刻原始结果。于是问题收敛为如何在不全量扫描的前提下快速定位这些高注意力 token。这正是 KTransformers CPU 稀疏注意力框架要解决的核心问题。3. 相关工作Prune 还是 Retrieval团队系统调研了近年的 KVCache 稀疏选择工作可以归为两大流派H2O 及后续工作Prune 流派最早指出推理注意力稀疏、只需保留 5% KVCache 的结论后续工作在此基础上设计更精细的 token 选择策略。这类方法对单条查询推理合理但如 Mooncake 项目的探索所指出的未来趋势是尽可能预计算可复用的 KVCache再用它回答不同问题“一次计算、多次复用”。因此 KTransformers 倾向于不删除 KVCache 中的 token或至少不删掉大部分保证不同问题可以各自关注上下文的不同部分。InfLLMRetrieval 流派其框架与上述理念高度契合被 KTransformers 选作核心参考。InfLLM 的核心设计将过长的上下文 KVCache 存入外部记忆模块Memory Units每步计算时只检索最相关的语义信息参与运算缓解“过长上下文导致注意力被无关噪声稀释”的问题外部记忆以语义块相邻 token 组成组织计算时采用滑动窗口机制——只选取上下文头部的块Initial Tokens、靠近当前 token 的块Local Tokens以及与当前 token 语义相似度最高的少量块参与注意力计算为高效检索InfLLM 在每个块内选取注意力得分 $$r_m$$ 最高的若干代表 token用相似度方程计算当前 token 与每个语义块的相关性见 InfLLM 公式图。与 H2O 相比InfLLM 的两个关键区别KVCache 不丢弃而是存放在内存中并在推理时动态加载到 GPUKVCache 以块为粒度管理每块选少量代表 token。其外部记忆单元可卸载到 CPU/DRAM 甚至 SSD按具体问题选取不同部分参与计算。3.1 Quest块级代表 token 与上界估计Quest 同样以块为粒度管理 token。它分析了 H2O 与全注意力下关键 token 的召回率发现 H2O 对 Top-10 注意力 token 的召回率仅约 50%关键信息丢失过多。为此 Quest 从每个块中选两个“代表 token”做检索Prefill 阶段每个 KVCache 块记录每个通道的最大值与最小值即图中标注的 Reduced Keys含逐元素 min key 与 max key注意力计算阶段当前 query 向量分别与每个块的 max key 和 min key 做点积对每个通道取两个乘积向量的最大值再求和作为该块相关性的上界根据相关性得分选出 top-k 块参与注意力计算。Quest 没有考虑异构架构假设全部 KVCache 仍可放进内存仅用稀疏注意力加速推理。最终 Quest 实现了 7.03x 注意力计算加速与 2.23x 端到端推理时延改善。3.2 SnapKV 及其他工作SnapKV 在 prefill 阶段一次性保留两部分 token图中橙色与绿色区段与 InfLLM 的区别只在于中间 token 的选择方式SnapKV 在token 级而非块级做选择打分方式与 H2O 类似即 $$softmax(\frac{qk^T}{\sqrt{d_k}})$$但按列求和时只统计最后绿色窗口对应 InfLLM 的 Local Tokens内的行并额外引入 pooling 操作论文解释是为了让召回 token 保留更完整的语义信息。SnapKV 是一次性选择选择完成后只用被选 token 做注意力其余 KVCache 直接丢弃。其他相关工作PyramidKV观察到注意力分数在层间呈金字塔分布——低层注意力分布广泛高层少数关键 token 的得分越来越突出。因此 PyramidKV 给低层分配更多 KVCache 容量、高层更少。MagicPiG基于 Locality-Sensitive HashingLSH的动态 KVCache 管理——先用 SnapKV 选一部分重要 token 存 GPU其余 KVCache 放内存利用 LSH 在高维空间检索中的高效性与 CPU 多线程能力检索与当前 query 相似的内存 KVCache 并载入。相比 InfLLM/Quest/SnapKVMagicPiG 不需要扫描全部代表 token 再选 top-k而是借助 LSH 的数学性质以低开销、高速度模拟注意力得分并定位重要 KVCache。4. KTransformers CPU 稀疏注意力框架4.1 设计要点从上述工作中KTransformers 提炼出四条关键结论注意力权重分布稀疏无用的 KVCache 可能引入噪声反而降低推理性能推理阶段 KVCache 驱逐策略的通行做法是保留 prompt 头尾的 token为中间部分设计选择算法模型性能的主要影响因素是能否准确识别关键 token中间部分按块管理可提升内存换入换出与注意力计算效率且小块粒度并不比 token 级表现更差各注意力层在推理时关注的 token 不同不同层分配到的 KVCache 容量也应不同。在此之上KTransformers 实现了通用的推理期稀疏 CPU 注意力算子框架。整体流程为Prefill 阶段采用 chunked prefill每次只把一层 KVCache 载入 GPU 计算算完存回 CPU/DRAMDecode 阶段不再来回换入换出 KVCache稀疏注意力算子直接运行在 CPU 上。这显著降低了最小 GPU 显存需求使本地 128K 甚至 1M token 上下文成为可能。生成阶段的完整框架如下4.2 KVCache 切分Initial / Context / LocalKVCache 按块组织完整输入 prompt 被切分为三个可配置部分Initial序列头部全程全量参与注意力Context中间部分稀疏检索Local序列尾部全程全量参与注意力。头尾全量保留的设计依据来自 streamingLLM、Minference 等论文发现的“注意力汇聚attention sinks”现象注意力权重往往显著偏向序列开头与结尾。4.3 Context 块划分与代表 token 选择中间 Context 部分沿用 InfLLM 思路按可配置固定 token 数划分为块每块可选 1k 个 token 作为代表 token推理时基于这些代表 token 挑选需要参与注意力的 Context 块。框架实现了五种代表 token 选择方法方法说明Max块内多个 token 逐通道取最大值拼成该块的 1 个代表 tokenMean块内多个 token 逐通道取均值拼成代表 tokenQuest取块内逐通道最大值与最小值作为代表 token代表 token 数固定为 2Dynamic按特定方法计算每个 token 的累计注意力得分每块选 top-k 最高分 token 作为代表类似 InfLLM 但做了简化Fix在块内按固定间隔选取 token代表 token 确定后使用 InfLLM 的相似度方程计算输入 X 与每个块 B 的 k 个代表 token 的相似度只选出 top $$r_k$$ 个块参与注意力计算其中 $$l_P$$ 为历史 token 长度。关于 InfLLM 原方法文档说明了两点工程取舍它需要在 prefill 阶段为每个 token 计算代表性得分、再依此选块代表 token这一操作对 prefill 实现是侵入式修改难以与其他方法整合且实测中多数场景下用其他方法组合即可达到类似甚至更好的效果。因此该方式最终未并入框架但Dynamic方法保留了其简化版。4.4 预选择Preselect借鉴 SnapKV 的可选模块框架初版对每个生成 token 都重新选择最相关的 KVCache 块。这既有明显的检索开销又发现频繁更换所选块并不会带来更好结果。例如 kvretrieval 数据集上模型回答常常前半正确、后半错误——由于答案是无意义的长字符串这说明前期 token 生成时选对了 KVCache 块后期却选错了。为此框架集成了 SnapKV 的方法推理开始前基于 question 的上下文 token 注意力得分预先选择一批相关 KVCache 块后续推理中块的选择被限制在这批预选范围内。在 kvretrieval 数据集上该机制可以 100% 选中包含正确答案的块。需要强调该方法严格依赖 Benchmark Prompt 的结构在复杂文档理解与生成等其他场景不保证最优因此被设计为可选模块。最终框架与可配置参数如下4.5 完整可配置参数一览threads_numCPU 线程数block_sizeKVCache 块大小local_windows_lenprompt 尾部窗口大小preselect_block_count预选块数量second_block_count预选后每次再选出的块数preselect_block是否启用预选择token_stepKVCache 块重选的 token 间隔layer_stepKVCache 块重选的层间隔dense_layer_num前若干层不做 KVCache 选择、全量参与head_select_modeSEPARATEGQA 场景下每个 kv_head 独立选择/SHARED所有 kv_head 联合选择representative_type配置文件名anchor_type代表 token 选择方法representative_num配置文件名anchor_num代表 token 数量通过调整这些参数可以复现多种 KVCache 驱逐/压缩方法例如block_size设为 1 且preselect_block为 True即得到无 pooling 的 SnapKV 版本representative_type设为Quest、preselect_block设为 False、head_select_mode设为SEPARATE即复现 Quest 方法。4.6 框架伪代码def preselect_block(local_q, kvcache): key_states kvcache.keycache attn_scores torch.matmul( local_q, key_states.transpose(2, 3) ) / math.sqrt(head_dim) attn_scores attn_mask attn_scores nn.functional.softmax( attn_scores, dim-1, dtypetorch.float32 ).to(query_states.dtype) vote attn_scores[..., initial_size:-local_size:, :].sum(dim-2) pool_vote pool1d(vote, kernel_sizekernel_size, paddingkernel_size//2, stride1) indices pool_vote.topk(max_capacity_prompt - local_size, dim-1).indices kv_cache_block_indices find_representative_tokens_block(indices) kvcache_after_preselected kvcache[kv_cache_block_indices] ... return kvcache_after_preselected def get_representative_tokens(): # 根据 representative_type 计算每个块的代表 token return ... def decode_attention(query, key, value): # 每 token_steps 个 token 选择一次 token_steps 4 # 每 layer_steps 层选择一次 layer_steps 4 for token_idx in range(max_new_tokens): for layer_idx in range(config.num_hidden_layers): if token_idx % token_steps ! 0 or layer_idx % layer_steps ! 0: # 本轮当前层无需重选时复用 kvcache 中保存的历史选择结果 kvcache_after_retrieval history_kvcache_after_retrieval[layer_idx//layer_steps] else: # 否则用当前轮当前层的 query 重新选择 kvcache kvcache_after_retrieval retrieval_kvcache(query, kvcache) # 保存进 kvcache 历史选择结果 history_kvcache_after_retrieval[layer_idx//layer_steps] kvcache_after_retrieval # 计算注意力 output attn(query, kvcache_after_retrieval) yield output # 模型 prefill若需预选择local_q 仍需保存 local_q, KVCache model.prefill(input_ids) if preselect_block: # 预选择轮 KVCache preselect_block(local_q, kvcache) # 找到每个块的代表 token block_representative_tokens get_representative_tokens( kvcache, config.representative_type ) # model generate decode_attention(query, key, value)伪代码体现了两个关键的工程优化token_steps/layer_steps让“选择”与“注意力计算”解耦——并非每个 token 的每层都重新检索而是按间隔重选、其余时间复用历史选择结果history_kvcache_after_retrieval按层组缓存直接降低了检索开销。5. 实验结果5.1 基线配置实验初始采用如下基础配置后续在扩展框架上进一步优化max_seq_len: 256000 # KVCache 长度 block_size: 128 # KVCache 块大小 local_windows_len: 4096 # 长度 local_windows_len 的 KVCache 存放在 GPU 上 second_block_count: 96 # 预选后每次选择的 KVCache 块数若 preselect_block_count则直接使用预选块 threads_num: 64 # CPU 线程数 representative_type: DYNAMIC # KVCache 块代表 token 选择方法 kv_type: FP16 dense_layer_num: 0 # 前若干层不填充/不选择 KVCache representative_num: 1 # 每个 KVCache 块内代表 token 数量 preselect_block: False # 是否启用预选择 head_select_mode: SHARED # 所有 kv_head 联合选择 preselect_block_count: 0 # 预选块数量 layer_step: 1 # 每几层选择一次 token_step: 1 # 每几个 token 选择一次在 opencompass 框架下passkey 数据集是在冗余文本的不同深度插入一小段数字kvretrieval 则是在随机生成的键值对中查找匹配项。128K 场景下原始模型与加速后 KTransformers 的对比Single needle retrieval zh 128k / passkey / kvretrieval方案Single needle retrieval zh 128kpasskeykvretrieval原始模型99.8910021.0KTransformers每个生成 token 重选 KVCache 块10010015.40简单数据集上两者均接近满分生成速度从 llama.cpp 的 4.86 tokens/s 提升到 KTransformers 的 27.49 tokens/s最高 5.65x 加速。较难的 kvretrieval 上初始配置明显掉分需要下一节的优化策略来弥补甚至超越原始模型。另外团队也用该配置框架复现了 Quest 的结果。由于 InternLM2.5-7B-Chat-1M 使用 GQAGrouped Query Attention而 Quest 论文主要针对 MHA 模型优化实测结果不理想Quest 官方也承认需要进一步支持 GQA故不再展开。5.2 预选择 稠密层超越原始模型前文提到逐 token 重选会导致文本越长越容易偏离、后期选到不同块。引入 SnapKV 式预选择后kvretrieval 分数从 15.4 提升到 24.2超过原始模型全注意力的 21 分。进一步发现 LLM 前几层 KVCache 的稀疏化效果不显著因此将前两层设为全量复用 KVCache最终取得24.4分相对原始 21.2 分。1M 数据集上的 needle-in-a-haystack 测试中KTransformers CPU 稀疏注意力框架不仅复现了原始模型报告分数还将准确率从 89.31 提升到92.88且推理速度达到 llama.cpp 的近 10 倍。5.3 与 llama.cpp 的系统级对比在 4090D 服务器上llama.cpp 配置为 KVCache 存 CPU/DRAM、全部计算在 GPU以 Single Needle Retrieval 数据集对比在保持100% 答题正确率的前提下prefill 提速 20.694.1 倍推理提速1.27.1 倍。prefill 差距大的主因llama.cpp 开启 KVCache offload 后在 CPU 上做注意力计算而长文本场景注意力不仅计算重还占据绝大部分计算时间KTransformers 则通过模板注入框架实现 GPU 逐层 Chunk Prefill。推理速度方面KTransformers 随 prompt 增长保持平稳近似水平线而 llama.cpp 随 prompt 变长逐渐变慢——因为 KTransformers 只选取最重要的 16K KVCache 块参与推理速度始终相当于 llama.cpp 处理 16K prompt 的水平不随上下文增长退化至少在这些测试数据集上如此。6. 使用指南适用范围长上下文功能目前仅通过local_chat.py接口支持仓库中对应实现位于 archive/ktransformers/local_chat.pyserver 接口的集成当时仍在开发中且long_context模式断言要求模型架构为LlamaForCausalLMInternLM2.5-1M 正是以 Llama 架构转换发布的。从 源码 可以看到约束与精度设置if mode long_context: assert config.architectures[0] LlamaForCausalLM, only LlamaForCausalLM support long_context mode torch.set_default_dtype(torch.float16)配套的优化规则文件为 Internlm2_5-7b-Chat-1m.yamlembed_tokens的 generate/prefill 均放在 CPULlamaModel与self_attn放在 CUDA 上执行——这正是“decode 在 GPU 做轻量计算、KVCache 主体驻留 CPU/DRAM”的算子级落点。团队已将对应的模型 config、GGUF 与 tokenizer 上传到公开的模型仓库按 InternLM2.5 转 Llama 1M 命名的仓库搜索即可获取。6.1 配置文件首次运行local_chat.py后会在~/.ktransformers下自动创建config.yaml由 Config 单例 从 仓库模板 拷贝而来。长上下文相关配置chunk_size: 4096 # prefill chunk size max_seq_len: 100000 # KVCache 长度 block_size: 128 # KVCache 块大小 local_windows_len: 4096 # 长度 local_windows_len 的 KVCache 存放在 GPU 上 second_select_num: 96 # 预选后每次选择的 KVCache 块数若 preselect_block_count则直接使用预选块 threads_num: 64 # CPU 线程数 anchor_type: DYNAMIC # KVCache 块代表 token 选择方法 kv_type: FP16 dense_layer_num: 0 # 前若干层不填充/不选择 KVCache anchor_num: 1 # 每个 KVCache 块内代表 token 数量 preselect_block: False # 是否启用预选择 head_select_mode: SHARED # 所有 kv_head 联合选择 preselect_block_count: 96 # 预选块数量 layer_step: 1 # 每几层选择一次 token_step: 1 # 每几个 token 选择一次注意源码中配置项命名为anchor_type/anchor_num/second_select_num见 config.py与文档概念名representative_type/representative_num一一对应。max_seq_len必须在~/.ktransformers/config.yaml中调整运行时 local_chat.py 会断言max_seq_len 输入长度 max_new_tokens。6.2 不同上下文长度所需内存上下文长度4K32K64K128K512K1MDRAM 占用 (GB)0.54.298.5817.168.7145.49请根据自身 DRAM 容量选择合适的max_seq_len。示例命令python local_chat.py --model_path/data/model/internlm2_5_to_llama_1m --gguf_path/data/model/internlm2_5_to_llama_1m --max_new_tokens500 --cpu_infer10 --use_cuda_graphTrue --modelong_context --prompt_file/path/to/file若已通过prompt_file指定了输入文本终端显示chat:后直接按回车即可开始推理。7. 小结KTransformers 的长上下文方案可以概括为三句话Prefill 阶段 GPU 逐层 chunk 计算、KVCache 落盘 CPU/DRAMDecode 阶段稀疏注意力直接跑在 CPU 上块级代表 token 间隔重选 可选的 SnapKV 预选择构成完整的选择策略。这一设计让单张 24GB GPU 可以在 128K 场景下以 7.1 倍于 llama.cpp 的速度、1M 场景下以约 16 tokens/sllama.cpp 近 10 倍的速度维持原生精度推理并且在 kvretrieval、needle-in-a-haystack 等基准上达到甚至超过原始模型全注意力的分数。其价值不仅在于具体的加速数字更在于提供了一个参数化、可复现 SnapKV/Quest 等多种方法的统一稀疏注意力框架为后续结合 MInference 等高性能稀疏 prefill 方法留出了扩展空间。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考