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

资讯详情

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

KTransformers 注入框架实战:用 YAML 规则将任意 Hugging Face 模型改造成 GPU/CPU 异构加速推理

KTransformers 注入框架实战:用 YAML 规则将任意 Hugging Face 模型改造成 GPU/CPU 异构加速推理 KTransformers 注入框架实战用 YAML 规则将任意 Hugging Face 模型改造成 GPU/CPU 异构加速推理【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformersKTransformers 是一个以 Python 为中心、可扩展的 LLM 推理优化框架其核心是一套基于模板的模块注入机制无需改动模型代码只需编写一个 YAML 规则文件并调用optimize_and_load_gguf即可把 Hugging Face Transformers 模型中的原始 torch 模块替换为 Marlin、Llamafile、自定义 MoE 等优化实现并利用 GPU/CPU 异构环境在资源受限的本地机器上运行大规模 MoE 模型。读完本文你将掌握注入框架的工作原理、YAML 规则的完整字段语义以及如何为一个自定义模型编写注入规则并完成量化权重加载与推理。本文内容以仓库中 archive/README_LEGACY.md 的 Quick Start 与 Brief Injection Tutorial 为主体并结合 注入框架实现、规则模板、local_chat 示例 等源码进行纵深讲解。背景KTransformers 是什么KTransformers读作 Quick Transformers的目标是增强 Transformers 的使用体验通过先进的 kernel 优化与“放置 / 并行”策略让研究者以极低的成本体验前沿 LLM 推理优化。框架定位是灵活、可扩展、以 Python 为中心用户只需一行代码注入一个优化模块即可获得 Transformers 兼容的接口、兼容 OpenAI 与 Ollama 的 RESTful API甚至一个简化的类 ChatGPT Web UI。KTransformers 特别关注本地部署场景。与 vLLM 侧重大规模部署优化不同KTransformers 面向资源受限的环境重视异构计算机会例如量化模型在 GPU/CPU 间的 offloadingCPU 侧使用高效的 Llamafile kernelGPU 侧使用 Marlin 4-bit kernel对应代码见 archive/ktransformers/ktransformers_ext/operators/custom_marlin。此外框架在演进过程中陆续支持了 Metax、SanechipsZhuFeng V1.0、Intel、Ascend、Kunpeng、AMD 等硬件厂商原文档 Quick Start 中明确列出的支持厂商并在 archive/ktransformers/util/vendors.py 中通过device_manager与GPUVendor统一抽象不同后端。核心概念模板化注入框架注入Injection是 KTransformers 的心脏。其设计目标非常明确让研究者轻松地用优化变体替换原始 torch 模块简化多种优化组合的过程从而探索它们之间的协同效应synergistic effects。整个流程是写一个 YAML 注入模板 → 在加载 Transformers 模型前调用optimize_and_load_gguf→ 框架递归遍历模型所有子模块按 YAML 中的规则做正则/类型匹配命中后替换为高级实现。最小使用示例原文档给出的核心代码片段如下with torch.device(meta): model AutoModelForCausalLM.from_config(config, trust_remote_codeTrue) optimize_and_load_gguf(model, optimize_config_path, gguf_path, config) ... generated prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens1000)流程拆解meta device 上构建模型with torch.device(meta)让模型只注册计算图与参数元数据不实际占用任何显存/内存为后续按规则逐模块加载权重做准备调用optimize_and_load_gguf该函数读取 YAML 规则遍历模型全部子模块按规则匹配并替换为高级模块同时从 GGUF 文件中加载对应张量推理注入完成后原始的generate接口仍然可用同时框架提供了兼容的prefill_and_generate方法可进一步启用 CUDA Graph 等优化来提升生成速度。prefill_and_generate的实现位于 archive/ktransformers/util/utils.py其签名如下def prefill_and_generate(model, tokenizer, inputs, max_new_tokens10000, use_cuda_graph: bool True, modenormal, force_think: bool False, chunk_size16384, use_flashinfer_mlaFalse, num_headsNone, head_dim_ckvNone, head_dim_kpeNone, q_head_dimNone, static_cacheNone, draft_modelNone, draft_cacheNone):从实现看该方法内部将流程拆为 prefill 阶段的chunk_prefill支持长上下文分块与 decode 阶段的单 token 循环当use_cuda_graphTrue时会通过 cuda_graph_runner 捕获 CUDA Graph 消除 kernel 启动开销若开启use_flashinfer_mla还会调用 flashinfer 的 MLA 规划接口MLAWrapperSingleton.plan_all加速 DeepSeek 系列模型的 MLA 注意力。YAML 注入规则match 与 replace每个规则由两个部分组成match声明要替换哪些模块。支持按模块名正则name和模块类class匹配两者可同时给出此时要求“名字与类同时命中”replace声明注入哪个模块类class以及初始化关键字参数kwargs。class可以是完整 Python 路径也可以写default表示保留原模块。原文档给出的经典示例——把所有torch.nn.Linear替换为 Marlin 4-bit 量化 kernel- match: name: ^model\\.layers\\..*$ # regular expression class: torch.nn.Linear # only match modules matching name and class simultaneously replace: class: ktransformers.operators.linear.KTransformersLinear # optimized Kernel on quantized data types device: cpu # which devices to load this module when initializing kwargs: generate_device: cuda generate_linear_type: QuantizedLinearMarlin注意原文档中replace使用了device字段与generate_linear_type参数这是早期版本v0.1.x的写法。在当前仓库的 规则模板 中字段已演进为generate_device/prefill_device与generate_op/prefill_op详见下文说明注入框架在迭代中把“生成阶段”与“预填充阶段”的算子选择解耦使 prefill 与 decode 可以分别使用不同 kernel——这正是异构优化的关键能力之一。规则引擎源码级拆解注入框架核心位于 archive/ktransformers/optimize/optimize.py包含三个关键函数。gen_optimize_config规则编译gen_optimize_configoptimize.py 第 67-118 行递归遍历模型树为每个模块名生成注入配置对每条规则依次检查matchclass匹配通过isinstance判断name匹配通过re.search判断若命中且该规则含replace则把class、kwargs记录到out_data[module_name]命中多条规则时按 YAML 顺序首个命中即break仅递归时继续多条规则对同一模块的kwargs会做合并out_data[module_name][kwargs].update(...)规则可带recursive字段控制是否继续递归注入子模块未命中任何规则的模块默认记为class: default并写入默认generate_device/prefill_device若match中class与name都缺失会抛出match must have at least one of \class\ and \name\异常——这是编写规则时最容易踩的坑。inject模块替换injectoptimize.py 第 28-54 行根据编译结果执行替换通过import_module_name/import_class_name动态导入replace.class指定的类将 GGUF 中对应key的张量设备映射记录到gguf_loader.tensor_device_map以module_cls(key..., gguf_loader..., config..., orig_modulechild, **kwargs)构造新模块并用set_module替换原模块对class default的模块只更新设备映射、保留原模块。optimize_and_load_gguf统一入口optimize_and_load_ggufoptimize.py 第 129-163 行是文档推荐的使用入口完整流程为读取并解析 YAML 规则文件调用gen_optimize_config编译注入配置调用translate_model_config处理特殊模型例如为 Mixtral 补全moe_intermediate_size在torch.device(meta)上下文内执行inject避免注入过程占用真实显存优先加载lm_head因其中间结果巨大再加载其余全部权重调用del_meta删除残留的 meta 参数并清理 CUDA / XPU 缓存。权重加载由 archive/ktransformers/util/custom_loader.py 中的GGUFLoader与SafeTensorLoader完成两者都实现了ModelLoader抽象基类提供has_tensor等接口因此注入模块可以统一通过gguf_loader句柄惰性加载张量。从规则模板看真实异构策略仓库 archive/ktransformers/optimize/optimize_rules 下存放了 DeepSeek-V2 / V3、Qwen2-57B-A14B、Mixtral、InternLM2.5、Moonlight、SmallThinker 等模型的完整规则模板它们是local_chat.py的默认优化配置。以 DeepSeek-V2-Lite-Chat.yaml 为例逐条解读- match: class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: cuda prefill_device: cuda把 DeepSeek-V2 的 YARN RoPE 替换为 KTransformers 的优化实现prefill 与 decode 都在 GPU 上执行。- match: name: ^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$ # regular expression class: torch.nn.Linear replace: class: ktransformers.operators.linear.KTransformersLinear kwargs: generate_device: cuda prefill_device: cuda generate_op: KLinearMarlin prefill_op: KLinearTorch这是最核心的一条用正则排除kv_b_projMLA 的 KV 压缩投影由专门实现处理把其余所有torch.nn.Linear替换为 KTransformersLinear。其中generate_op: KLinearMarlin表示 decode 阶段使用 GPU Marlin 4-bit kernelprefill_op: KLinearTorch表示 prefill 阶段退回 PyTorch 原生算子——两类算子分离是兼容性与性能的折中。- match: name: ^lm_head class: torch.nn.Linear replace: class: ktransformers.operators.linear.KTransformersLinear kwargs: generate_device: cuda prefill_device: cuda generate_op: KLinearMarlin prefill_op: KLinearTorchlm_head单独一条规则因为它的输入是中间隐藏状态、输出词表很大值得单独指定算子。- match: name: ^model\\.layers\\..*\\.mlp$ class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE kwargs: generate_device: cuda prefill_device: cuda - match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts kwargs: prefill_device: cuda prefill_op: KExpertsTorch generate_device: cpu generate_op: KExpertsCPU out_device: cuda recursive: False # dont recursively inject submodules of this module这两条展示了 KTransformers 最具代表性的专家 offloading 策略MoE 模块本身mlp在 GPU 上处理路由逻辑而专家权重mlp.experts在prefill 阶段放 GPU、decode 阶段放 CPUgenerate_device: cpuout_device: cuda把专家计算结果送回 GPU 与主分支汇合。recursive: False防止继续注入专家的内部子模块。decode 阶段使用KExpertsCPU即 Llamafile 系列的 CPU kernel 完成量化专家计算这正是“GPU/CPU 混合推理 MoE”的核心思想对应论文标题KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models。- match: name: ^model\\.layers\\..*\\.self_attn$ replace: class: ktransformers.operators.attention.KDeepseekV2Attention kwargs: generate_device: cuda prefill_device: cudaMLA 注意力整体替换为优化实现KDeepseekV2Attention。- match: name: ^model$ replace: class: ktransformers.operators.models.KDeepseekV2Model kwargs: per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill - match: name: ^model.embed_tokens replace: class: default kwargs: generate_device: cpu prefill_device: cpu模型主干替换为KDeepseekV2Model支持 layer-wise prefill 阈值控制而 embedding 层保留原实现class: default并把设备放到 CPU避免占用宝贵的显存。规则文件命名约定与默认配置在local_chat.pyarchive/ktransformers/local_chat.py中框架为不同模型架构内置了默认规则文件模型架构类默认规则文件DeepseekV2ForCausalLMDeepSeek-V2-Chat.yamlDeepseekV3ForCausalLMDeepSeek-V3-Chat.yamlQwen2MoeForCausalLMQwen2-57B-A14B-Instruct.yamlLlamaForCausalLMInternlm2_5-7b-Chat-1m.yamlMixtralForCausalLMMixtral.yaml当你运行local_chat且不指定--optimize_config_path时框架会根据config.architectures[0]自动挑选对应规则文件只有模型不在内置列表时才会交互式询问规则文件路径。规则文件所在目录还按硬件平台组织子目录npu/、rocm/、xpu/对应 Ascend NPU、AMD ROCm、Intel Arc 等平台的特化规则。从示例到自定义模型编写自己的规则综合原文档与源码为一个新模型编写注入规则可遵循以下步骤分析模型结构用AutoConfig.from_pretrained加载配置查看num_hidden_layers、architectures等字段确定需要替换的模块层级model.layers.*.self_attn、mlp、lm_head等确定替换目标线性层 →KTransformersLinearMoE →KDeepseekV2MoE/KTransformersExpertsMLA 注意力 →KDeepseekV2AttentionRoPE → 对应RoPE算子见 archive/ktransformers/operators/RoPE.py编写规则注意name是相对模型根的正则表达式须转义点号\\.class匹配依赖isinstance因此要写模块实际类含完整包路径不要忘记为未命中规则的模块保留默认分支或加recursive: False控制注入深度验证在 meta device 上构建模型后调用optimize_and_load_gguf观察终端打印的Injecting module as class日志确认每个预期模块都被替换且 GGUF 张量名称能正确对应名称翻译逻辑见 archive/ktransformers/util/custom_gguf.py 的translate_name_to_gguf。附引用与致谢若在研究中使用了 KTransformers原文档提供了 BibTeX 引用条目该论文发表于 ACM SIGOPS 2025inproceedings{10.1145/3731569.3764843, title {KTransformers: Unleashing the Full Potential of CPU/GPU Hybrid Inference for MoE Models}, author {Chen, Hongtao and Xie, Weiyu and Zhang, Boxin and Tang, Jingqi and Wang, Jiahao and Dong, Jianwei and Chen, Shaoyuan and Yuan, Ziwei and Lin, Chen and Qiu, Chengyu and Zhu, Yuening and Ou, Qingliang and Liao, Jiaqi and Chen, Xianglin and Ai, Zhiyuan and Wu, Yongwei and Zhang, Mingxing}, booktitle {Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles}, year {2025} }KTransformers 基于 Transformers 框架开发并受益于 GGUF/GGML、Llamafile、Marlin、SGLang、flashinfer 等开源生态更多常见问题可查阅 doc/en/FAQ.md模型级注入与多 GPU 的详细教程可参考 doc/en/injection_tutorial.md 与 doc/en/deepseek-v2-injection.md。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表