尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

kTransformers 长上下文推理实战:local_chat 模式下的 1M Token KVCache 管理方案

kTransformers 长上下文推理实战:local_chat 模式下的 1M Token KVCache 管理方案 kTransformers 长上下文推理实战local_chat 模式下的 1M Token KVCache 管理方案【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers本文基于 kTransformers 仓库中的长上下文使用教程doc/en/long_context_tutorial.md系统讲解如何通过local_chat.py接口启用long_context模式在 24G 显存的单卡机器上运行 InternLM2.5-7B-Chat-1M 的百万 token 级长上下文推理。读完本文你将掌握长上下文模式下的完整启动流程、~/.ktransformers/config.yaml中全部 KVCache 管理参数的含义与默认值、按 DRAM 容量规划max_seq_len的方法以及分块预填充chunk prefill与 CPU 卸载等底层实现细节。一、适用场景与当前能力边界kTransformers 的长上下文框架面向“输入序列远超 GPU 显存可承载 KVCache 容量”的场景把 KVCache 按块block拆分热块保留在 GPU、冷块下沉到 CPU 内存DRAM并在每一层通过锚点anchor评分动态选回最相关的 KVCache 块。仓库文档明确说明了当前的能力边界见 doc/en/long_context_tutorial.md长上下文目前仅由local_chat.py接口支持server 接口的集成尚在开发中官方给出的参考模型是InternLM2.5-7B-Chat-1M1M 上下文能力在24G 显存的 GPU 上即可运行从源码看long_context模式对模型架构有硬性限制。archive/ktransformers/local_chat.py 中存在如下断言if mode long_context: assert config.architectures[0] LlamaForCausalLM, only LlamaForCausalLM support long_context mode torch.set_default_dtype(torch.float16)也就是说model_path指向的模型architectures必须是LlamaForCausalLMInternLM2.5 系列模型正是基于 LLaMA 架构改造因此满足该约束并且推理 dtype 会被强制切换为float16。这一约束与教程中“仅支持 Llama 架构模型”的用法完全对应。二、准备模型资源model_path 与 gguf_path为了方便用户管理长上下文所需的资源官方将模型 config、GGUF 量化权重、tokenizer打包上传到了一个模型仓库仓库名nilv234/internlm2_5_to_llama_1m托管于 Hugging Face可按下述仓库名在 Hugging Face 上检索获取。使用方式上只需把local_chat函数的两个参数指向同一个目录model_path/path/to/repo—— 提供config.json、tokenizer等 Hugging Face 格式文件gguf_path/path/to/repo—— 提供 CPU 侧加载的 GGUF 量化权重。local_chat的完整参数签名如下见 archive/ktransformers/local_chat.py其中与长上下文直接相关的参数是mode、chunk_size、prompt_filedef local_chat( model_path: str | None None, optimize_config_path: str None, gguf_path: str | None None, max_new_tokens: int 1000, cpu_infer: int Config().cpu_infer, use_cuda_graph: bool True, prompt_file : str | None None, mode: str normal, force_think: bool False, chunk_size: int 8192, device: str cuda, tp: int 1, ):对LlamaForCausalLM架构优化规则会自动命中默认规则文件 archive/ktransformers/optimize/optimize_rules/Internlm2_5-7b-Chat-1m.yamllocal_chat.py 中的default_optimize_rules映射无需手动指定optimize_config_path。三、启动命令在仓库的ktransformers目录下执行长上下文相关代码位于 archive/ktransformers/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各参数含义参数作用--model_path/--gguf_path分别指定 HF 格式模型目录与 GGUF 权重目录长上下文场景下两者指向同一仓库目录--max_new_tokens本次生成最多产出的 token 数--cpu_infer分配给 CPU 侧算子GGUF 层的线程数教程示例为 10--use_cuda_graph对 decode 阶段使用 CUDA Graph 加速--mode设为long_context启用长上下文框架默认值normal--prompt_file预置输入文本文件已在该文件给出输入时终端出现chat:提示后直接按回车即可开始推理关于prompt_file的交互逻辑可以从 local_chat.py 确认主循环读取到空输入直接按回车时若提供了prompt_file则读取该文件内容作为本轮输入否则视为无操作继续等待。因此“文件给出输入 按回车”是教程推荐的免手敲长文本方式。四、配置文件机制~/.ktransformers/config.yaml的自动创建长上下文的绝大多数参数不在命令行而是集中在一个 YAML 配置中。教程指出首次运行local_chat.py后会在~/.ktransformers下自动创建config.yaml文件。这一机制在 archive/ktransformers/server/config/config.py 的Config.load()中有明确实现以ktransformers/configs/config.yaml为模板见 archive/ktransformers/configs/config.yaml若用户目录~/.ktransformers不存在则创建若~/.ktransformers/config.yaml尚不存在则把仓库内模板复制过去之后所有运行都加载并解析用户目录下的这份文件。这意味着首次运行后请编辑~/.ktransformers/config.yaml而不是仓库内的模板修改才会对你生效。Config是单例类long_context段的解析见 config.py。五、long_context参数全解模板配置中long_context段位于 archive/ktransformers/configs/config.yaml。结合教程给出的注释与Config中的默认值各参数说明如下long_context: chunk_size: 4096 # prefill 分块大小 max_seq_len: 100000 # KVCache 总长度token 数 block_size: 128 # KVCache 块大小 local_windows_len: 4096 # 常驻 GPU 的 KVCache 窗口长度 second_select_num: 96 # 粗筛后每轮选回的 KVCache 块数若 preselect_block_count 则直接使用预选块 threads_num: 64 # CPU 线程数 anchor_type: DYNAMIC # KVCache 块代表 tokenanchor的选取方式 kv_type: FP16 # KVCache 存储精度 dense_layer_num: 0 # 前若干层不参与 KVCache 填充/选择 anchor_num: 1 # 每个 KVCache 块内的代表 token 数 preselect_block: False # 是否开启块预选 head_select_mode: SHARED # 所有 kv_heads 联合参与选择 preselect_block_count: 96 # 预选块数量 layer_step: 1 # 每隔几层执行一次选择 token_step: 1 # 每隔几个 token 执行一次选择逐项展开chunk_size默认 4096prefill 阶段的分块大小。长上下文中 prefill 不可能一次性把整段输入送进模型而是按 chunk 循环处理。archive/ktransformers/util/utils.py 的prefill_wrapper展示了该循环def prefill_wrapper(profNone): nonlocal logits chunk_start 0 while chunk_start seq_length: chunk_end min(chunk_start chunk_size, seq_length) if past_key_values ! None: past_key_values.cur_idx cache_position[chunk_start:chunk_end] logits chunk_prefill(inputs[:, chunk_start:chunk_end], cache_position[chunk_start:chunk_end], past_key_values) chunk_start chunk_size注意local_chat命令行还有一个同名参数chunk_size默认 8192它会写入config.chunk_size并透传给prefill_and_generate与 KVCache 块管理的chunk_size是不同层面的控制量。max_seq_len默认 32000教程示例 100000KVCache 池可容纳的最大 token 数是显存/内存占用的第一决策参数需结合 DRAM 容量选择见第六节。local_chat.py在每轮推理前会做校验local_chat.pyif mode long_context: assert Config().long_context_config[max_seq_len] input_tensor.shape[1] max_new_tokens, \ please change max_seq_len in ~/.ktransformers/config.yaml即输入 token 数加max_new_tokens不得超过max_seq_len否则报错提示修改~/.ktransformers/config.yaml。block_size默认 128KVCache 的最小管理单元是“块”每块覆盖 128 个 token 位置的 K/V 张量。块是“存到 DRAM / 选回 GPU”的搬运粒度。local_windows_len默认 4096常驻 GPU 的 KVCache 窗口长度。以block_size128计约相当于 32 个块的热窗口是显存预算与召回质量的直接调节项。second_select_num默认 32教程示例 96每层在粗筛preselect之后进一步选回的 KVCache 块数量当它大于等于preselect_block_count时可以直接复用预选块而省去二次评分。anchor_type: DYNAMIC默认 DYNAMIC/anchor_num默认 1为每个 KVCache 块挑选 1 个代表性 token 作为锚点用锚点的注意力得分近似整个块的重要性。DYNAMIC表示锚点位置随内容动态选取而非固定块内首/尾 token。kv_type: FP16默认 FP16KVCache 的存储精度。KVCache 是长上下文场景下的存储大头精度选择直接影响 DRAM 占用与召回质量。dense_layer_num默认 2最前面的若干层不做 KVCache 填充与选择全量保留保证浅层注意力不受块丢弃影响。教程示例中将其设为 0即所有层均参与块管理。preselect_block默认 True/preselect_block_count默认 32两级选择的第一级——按层/头粒度粗筛出一批候选块把每层的细粒度选择成本从“全部块”降到“候选块”。教程示例将preselect_block关闭False。head_select_mode: SHARED默认 SHARED多个 KV head 联合共享做块选择而非各 head 独立选择从而降低选择计算开销并统一块的驻留状态。layer_step默认 1/token_step默认 100教程示例 1选择操作的频率控制——每隔layer_step层、每生成token_step个 token 才重新执行一次锚点评分与块选回。token_step1表示逐 token 都做选择最精细也最开销大调大可以摊销选择成本。threads_num教程示例 64CPU 侧 KVCache 搬运与评分使用的线程数可按物理核数调整。注教程中的示例值如max_seq_len: 100000、second_select_num: 96、preselect_block: False与模板默认值32000、32、True并不完全一致前者是为 1M 级输入实测后推荐的配置可按自己的 DRAM/显存条件在两者之间取平衡。六、内存规划按 DRAM 容量选择max_seq_len教程给出了 InternLM2.5-7B-Chat-1M 场景下不同上下文长度对应的 DRAM 占用KVCache 主体存放在内存DRAM Size (GB)4K32K64K128K512K1M4K0.532K4.2964K8.58128K17.1512K68.71M145.49整理为便于查表的格式上下文长度4K32K64K128K512K1MDRAM 占用 (GB)0.54.298.5817.168.7145.49可见占用大致随长度线性增长128K 约 17.1 GB1M 约 145.49 GB。因此64 GB 内存的机器max_seq_len建议控制在 64K 附近128 GB 内存可支撑 128K256K跑满 1M 上下文至少需要准备 150 GB 量级的 DRAM。选择时应留出模型权重GGUF 常驻内存与其他进程的余量避免按上表“顶格”配置导致系统内存不足。七、底层实现速览KVCache 为什么能放下 1M token结合源码长上下文模式的执行路径与普通模式有三处关键差异都在 archive/ktransformers/util/utils.py 的prefill_and_generate中不使用 GPU 侧 StaticCache。普通模式下past_key_values会创建为一个 GPU 上的StaticCache容量seq_length max_new_tokens而long_context模式下该对象直接置为Noneutils.pyelif mode ! long_context: past_key_values StaticCache( config model.config, max_batch_size 1, max_cache_len seq_length max_new_tokens, device device_map, dtype model.dtype ) else: past_key_values None缓存的管理权交给模型自定义的长上下文 KVCache 后端按block_size分块、GPU 热窗口 DRAM 冷存储、按锚点评分选块这正是“显存固定、内存扩容”的实现基础。Embedding 在 CPU 上完成。分块预填充函数中long_context模式下 token id 到 embedding 的转换发生在 CPU避免把超大输入序列直接堆到显存utils.pydef chunk_prefill(inputs, cache_position, past_key_values): if mode long_context: inputs_embeds model.model.embed_tokens(inputs.to(cpu)) else: inputs_embeds model.model.embed_tokens(inputs.to(cpu)).to(torch_device)配合local_chat.py中LlamaForCausalLM被强制eagerattention 实现local_chat.py使得逐块注意力可以走自定义的选块内核。分块循环 逐层选块。prefill 按chunk_size推进每层在计算注意力前依据锚点anchor_type/anchor_num对 KVCache 块评分选回second_select_num个块进入local_windows_len热窗口preselect_block开启时先做一次粗筛layer_step/token_step控制重评分频率。这些行为全部由第五节的参数驱动。八、操作要点与常见问题只按回车就能跑使用--prompt_file预置输入后终端显示Chat:提示时直接按回车即可不必再粘贴文本报错please change max_seq_len in ~/.ktransformers/config.yaml说明输入长度加max_new_tokens超过了max_seq_len调大该值并确认 DRAM 余量对照第六节表格报错only LlamaForCausalLM support long_context mode当前model_path的config.json中architectures不是LlamaForCausalLM长上下文模式暂不支持该架构改配置不生效确认修改的是~/.ktransformers/config.yaml而非仓库内 configs/config.yaml 模板——模板只在首次运行时被复制一次显存不足优先下调max_seq_len其次减小local_windows_lenGPU 热窗口与chunk_sizeprefill 分块吞吐偏低尝试增大layer_step、token_step以降低选块频率或开启preselect_block利用粗筛复用同时按 CPU 核数调整threads_num与--cpu_infer。小结kTransformers 的长上下文框架通过local_chat.py --modelong_context提供了“24G 显存 大容量 DRAM”组合下运行 InternLM2.5-7B-Chat-1M 级别 1M 上下文推理的完整方案KVCache 按 128 token 的块管理热窗口驻留 GPU、冷块下沉 DRAM锚点评分决定每层选回哪些块所有调优旋钮集中在~/.ktransformers/config.yaml的long_context段配合第六节的 DRAM 占用表即可按机器内存精确设定max_seq_len。由于 server 接口集成仍在开发中生产化部署请持续关注 doc/en/long_context_tutorial.md 的后续更新。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表