企业官网建设流程全解析

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

【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers

KTransformers 是一个以 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 间的 offloading：CPU 侧使用高效的 Llamafile kernel，GPU 侧使用 Marlin 4-bit kernel（对应代码见 archive/ktransformers/ktransformers_ext/operators/custom_marlin）。此外，框架在演进过程中陆续支持了 Metax、Sanechips（ZhuFeng 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_code=True) optimize_and_load_gguf(model, optimize_config_path, gguf_path, config) ... generated = prefill_and_generate(model, tokenizer, input_tensor.cuda(), max_new_tokens=1000)

流程拆解：

1. meta device 上构建模型：with torch.device("meta")让模型只注册计算图与参数元数据，不实际占用任何显存/内存，为后续按规则逐模块加载权重做准备；
2. 调用optimize_and_load_gguf：该函数读取 YAML 规则，遍历模型全部子模块，按规则匹配并替换为高级模块，同时从 GGUF 文件中加载对应张量；
3. 推理：注入完成后，原始的generate接口仍然可用；同时框架提供了兼容的prefill_and_generate方法，可进一步启用 CUDA Graph 等优化来提升生成速度。

prefill_and_generate的实现位于 archive/ktransformers/util/utils.py，其签名如下：

def prefill_and_generate(model, tokenizer, inputs, max_new_tokens=10000, use_cuda_graph: bool = True, mode='normal', force_think: bool = False, chunk_size=16384, use_flashinfer_mla=False, num_heads=None, head_dim_ckv=None, head_dim_kpe=None, q_head_dim=None, static_cache=None, draft_model=None, draft_cache=None):

从实现看，该方法内部将流程拆为 prefill 阶段的chunk_prefill（支持长上下文分块）与 decode 阶段的单 token 循环，当use_cuda_graph=True时会通过 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_config（optimize.py 第 67-118 行）递归遍历模型树，为每个模块名生成注入配置：

• 对每条规则依次检查match：class匹配通过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：模块替换

inject（optimize.py 第 28-54 行）根据编译结果执行替换：

1. 通过import_module_name/import_class_name动态导入replace.class指定的类；
2. 将 GGUF 中对应key的张量设备映射记录到gguf_loader.tensor_device_map；
3. 以module_cls(key=..., gguf_loader=..., config=..., orig_module=child, **kwargs)构造新模块，并用set_module替换原模块；
4. 对class == "default"的模块只更新设备映射、保留原模块。

optimize_and_load_gguf：统一入口

optimize_and_load_gguf（optimize.py 第 129-163 行）是文档推荐的使用入口，完整流程为：

1. 读取并解析 YAML 规则文件；
2. 调用gen_optimize_config编译注入配置；
3. 调用translate_model_config处理特殊模型（例如为 Mixtral 补全moe_intermediate_size）；
4. 在torch.device("meta")上下文内执行inject，避免注入过程占用真实显存；
5. 优先加载lm_head（因其中间结果巨大），再加载其余全部权重；
6. 调用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_proj（MLA 的 KV 压缩投影，由专门实现处理），把其余所有torch.nn.Linear替换为 KTransformersLinear。其中generate_op: "KLinearMarlin"表示 decode 阶段使用 GPU Marlin 4-bit kernel，prefill_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: "KLinearTorch"

lm_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 # don't recursively inject submodules of this module

这两条展示了 KTransformers 最具代表性的专家 offloading 策略：MoE 模块本身（mlp）在 GPU 上处理路由逻辑，而专家权重（mlp.experts）在prefill 阶段放 GPU、decode 阶段放 CPU（generate_device: "cpu"），out_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: "cuda"

MLA 注意力整体替换为优化实现（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.py（archive/ktransformers/local_chat.py）中，框架为不同模型架构内置了默认规则文件：

当你运行local_chat且不指定--optimize_config_path时，框架会根据config.architectures[0]自动挑选对应规则文件；只有模型不在内置列表时才会交互式询问规则文件路径。规则文件所在目录还按硬件平台组织子目录（npu/、rocm/、xpu/），对应 Ascend NPU、AMD ROCm、Intel Arc 等平台的特化规则。

从示例到自定义模型：编写自己的规则

综合原文档与源码，为一个新模型编写注入规则可遵循以下步骤：

1. 分析模型结构：用AutoConfig.from_pretrained加载配置，查看num_hidden_layers、architectures等字段，确定需要替换的模块层级（model.layers.*.self_attn、mlp、lm_head等）；
2. 确定替换目标：线性层 →KTransformersLinear；MoE →KDeepseekV2MoE/KTransformersExperts；MLA 注意力 →KDeepseekV2Attention；RoPE → 对应RoPE算子（见 archive/ktransformers/operators/RoPE.py）；
3. 编写规则：注意name是相对模型根的正则表达式，须转义点号（\\.）；class匹配依赖isinstance，因此要写模块实际类（含完整包路径）；不要忘记为未命中规则的模块保留默认分支或加recursive: False控制注入深度；
4. 验证：在 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 2025）：

@inproceedings{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