KTransformers 操作符注入教程以 DeepSeek-V2 为例逐层替换自定义算子【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers导读KTransformers 的 Inject 注入框架允许用户在不修改 transformers 源码的前提下通过 YAML 规则将模型中的任意子模块替换为自定义算子从而实现 MLA 注意力、量化线性层、CPU 专家并行等异构优化。本教程以 DeepSeekV2-Chat 为例从注入规则语法出发依次演示 MLA 注入、路由专家注入、线性层注入、RoPE 缓冲区补偿、设备指定与多 GPU 划分最后讲解如何基于BaseInjectedModule编写一个可注入的全新算子。读完本文你将能够独立为任意 transformers 模型编写注入规则并把自研算子挂载进 KTransformers 的推理链路。How to Write Injection Rules注入规则的基础语法注入规则本质是一个 YAML 列表KTransformers 的注入框架在模型初始化阶段逐条读取规则对模型树上的每个模块做正则/类型匹配命中后以替换类重建该模块。规则基本形式如下- match: name: ^model\\.layers\\..*\\.*$ # Target module name class: torch.nn.Linear # Target module replace: class: default kwargs: generate_device: cuda:0 # your_op_param_1: 1234 # your_op_param_2: 5678 recursive: True各字段含义match匹配规则包含name与class两种形式。两者可以同时出现也可以单独出现同时出现时要求模块名称与类型均命中才会注入。name正则表达式匹配模块在模型中的完整路径名如model.layers.0.self_attnclass目标模块的 Python 类例如torch.nn.Linear、ktransformers.models.modeling_deepseek.DeepseekV2MoE。replaceclass用于替换目标模块的 Python 类路径。若不想替换仅设置设备等属性设为defaultkwargs替换模块初始化所需的参数字典其中generate_device指定该模块运行的设备可取值cpu、cuda、cuda:1等。recursive是否递归处理该模块的子模块默认True。recursive字段非常关键像 Self-attention 这类模块内部通常包含 q/k/v/o 四个线性子模块如果替换了整个自注意力模块却不想让内部线性层再被其他规则覆盖就需要把该规则设为False典型场景见下文 Routed Experts 注入。从实现看optimize.py 中的gen_optimize_config会深度优先遍历模型树对每个模块依次尝试规则列表class匹配通过isinstance判断name匹配通过re.search判断首个命中的规则生效并break随后inject同文件 L28-L54动态import替换类用module_cls(key..., gguf_loader..., config..., orig_modulechild, **kwargs)构造新模块并通过set_module挂回原位置。因此规则列表中规则的书写顺序即匹配优先级先写的规则优先生效。Understanding Model Structure从权重文件反推模型结构KTransformers 注入框架自由度很高允许你替换/实验任何基础算子但前提是必须清楚所运行模型的内部结构。以deepseek-ai/DeepSeek-V2-Lite-Chat为例打开模型仓库文件列表从.safetensors权重文件名可以看到每一层权重的命名对应注入规则中的match.name从modeling_deepseek.py文件可以看到每个模块类的具体实现对应注入规则中的match.class。由此可得到 DeepSeekV2 的模型骨架model.embed_tokens词嵌入→ 60 层model.layers.N每层含self_attn与mlpmlp内部是gate/experts→model.norm→lm_head。上述self_attn、mlp.experts、lm_head等正是注入规则中正则要命中的路径。KTransformers 官方支持的算子与后端对照表如下对应仓库内ktransformers/operators下的实现matchreplacebackendsdescriptionsLinearKTransformersLinearKLinearMarlinMarlin as backendKLinearTorchpytorch as backendKLinearCPUInferllamafile as backendKLinearFP8Triton fp8_gemm kernel. Requires GPU be able to calculate fp8 dataexpertsKTransformersExpertsKExpertsTorchpytorch as backendKExpertsMarlinMarlin as backendKExpertsCPUllamafile as backendAttentionKDeepseekV2AttentionKDeepseekV2AttentionMLA implementationMoEKMistralSparseMoEBlockKQwen2MoeSparseMoeBlockMoE for Qwen2KDeepseekV2MoEKDeepseekV2MoEMoE for DeepseekV2ModelKQwen2MoeModelKQwen2MoeModelModel for Qwen2KDeepseekV2ModelKDeepseekV2ModelModel for DeepseekV2RoPERotaryEmbeddingRotaryEmbeddingRoPE moduleYarnRotaryEmbeddingYarnRotaryEmbeddingRoPE module在 archive/ktransformers/operators/attention.py 中可以看到KDeepseekV2Attention以(BaseInjectedModule, DeepseekV2Attention)双重继承的方式实现experts.py 中KTransformersExperts通过EXPERTS_MAP[generate_op]分发到KExpertsCPU/KExpertsTorch等具体后端linear.py 中KTransformersLinear同样以generate_op选择KLinearMarlin等实现。generate_op就是你在 YAML 里指定具体后端的关键字。下面以 DeepSeekV2-Chat 为目标逐步注入自定义 Marlin 线性层、基于 Absorption 的 MLA 注意力、自定义专家模块、自定义 MoE 模块、自定义 RoPE 模块并为每个模块指定运行设备。Matrix Absorption-based MLA Injection替换注意力为 MLA 实现MLAMulti-head Latent Attention通过矩阵吸收把 q/k/v 投影合并显著降低 KV Cache 占用。注入 Attention 模块只需用正则命中 transformers 中的模块名替换为 KTransformers 的 MLA 实现- match: name: ^model\\.layers\\..*\\.self_attn$ # Regular expression replace: class: ktransformers.operators.attention.KDeepseekV2Attention # Optimized MLA implementation规则中match负责指定被替换模块replace指定注入的模块类及初始化关键字。完整的DeepSeek-V2-Chat.yaml中还为该规则补充了设备参数- match: name: ^model\\.layers\\..*\\.self_attn$ replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: cuda prefill_device: cudaprefill_device与generate_device分别控制预填充prefill阶段与逐 token 生成generate阶段的运行设备两阶段可以异构这是 KTransformers 异构推理的核心设计之一见 base_operator.py 中BaseInjectedModule对这两个属性的登记。Injection of Routed Experts注入路由专家并改写 MoE 前向对于图中标注为exps的路由专家Routed Experts注入的是被包装模块KTransformersExperts包裹的 CPUInfer 算子。KTransformersExperts有多个后端实现需要通过关键字告诉包装模块选择哪个实现以及如何使用它。transformers 中 MoE 用nn.ModuleList实现专家列表我们不希望 KTransformers 遍历列表逐个注入子模块因此在规则中设置recursive: False- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # Custom MoE kernel with expert parallelism kwargs: generate_device: cpu generate_op: MLPCPUExperts out_device: cuda recursive: False # Dont recursively inject submodules of this modulegenerate_device: cpu专家权重与计算放在 CPU由 llamafile/CPUInfer 承载generate_op: MLPCPUExperts选择 CPU 专家后端out_device: cuda专家输出写回 GPU供后续算子使用。实际仓库中的DeepSeek-V2-Chat.yaml同时为 prefill 阶段配置了prefill_op: KExpertsTorch、prefill_device: cuda即 prefill 用 GPU 上的 PyTorch 专家、generate 用 CPU 专家实现两阶段后端分离- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism 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注意官方 YAML 中 generate 阶段用的后端名是KExpertsCPU即文档正文中的MLPCPUExperts演变而来对应 experts.py 中out_device参数“输出应落在哪个设备”的语义——代码中KExpertsCPU.output_gpu_map[out_device]会为每个目标 GPU 预分配输出缓冲并通过submit_with_cuda_stream与 CUDA 流异步协同。注入路由专家为自定义模块后原nn.ModuleList的接口不再可用因此必须修改 FFN 模块的forward函数。最简单的做法是实现一个带自定义forward的新模块并注入- match: class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # MLP module with custom forward functionKDeepseekV2MoE定义在 archive/ktransformers/operators/experts.py它重写了 MoE 的前向逻辑以适配被替换后的专家模块调用方式。Injection of Linear Layers用类型检查精确注入量化线性层对于剩余线性层我们希望使用量化算子来节省存储并提升性能。由于目前没有 MLA 与量化联合使用的研究我们不希望把线性层注入到 MLA 算子内部因此通过正则负向断言 类型检查双重约束只有 name 与 class同时命中的模块才会被注入- match: name: ^model\\.layers\\.(?!.*self_attn).*$ # 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 kwargs: generate_device: cuda generate_op: QuantizedLinearMarlin(?!.*self_attn)是正则负向前瞻排除所有包含self_attn的路径从而保证 MLA 模块内部不被线性层规则波及。官方DeepSeek-V2-Chat.yaml中该规则更精确地排除了kv_b_projMLA 吸收合并后不再单独需要它并分别指定了生成与预填充阶段的后端- match: name: ^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$ # 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 kwargs: generate_device: cuda prefill_device: cuda generate_op: KLinearMarlin prefill_op: KLinearTorch即生成阶段使用 Marlin 量化内核KLinearMarlin预填充阶段退回 PyTorch 实现KLinearTorch。KLinearMarlin封装了 GPTQ 量化权重与 Marlin 反量化内核linear.py 的KLinearBase基类会在初始化时从orig_module或 GGUF 权重元信息中解析in_features/out_features并维护loaded状态。Injection of Modules with Pre-calculated Buffers补偿 RoPE 的预计算缓冲区为避免初始化注入后的原模型占用资源KTransformers 使用torch.device(meta)元设备初始化原模型见 optimize.pywith torch.device(meta): inject(...)随后再load_weights并del_meta清理 meta 参数。但 RoPE 模块在初始化时会预计算一批缓冲区如频率表元设备下不会真正执行计算因此需要在加载模型时补偿这些缓冲区的计算。做法是向旋转嵌入模块注入自定义模块在加载阶段完成预计算- match: class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbeddingktransformers.operators.RoPE.YarnRotaryEmbedding定义于 archive/ktransformers/operators/RoPE.py继承了原DeepseekV2YarnRotaryEmbedding的语义但在实际设备上完成缓冲区预计算从而让基于 YARN 的长上下文旋转位置编码在注入框架下正常工作。官方规则同时为其指定了generate_device: cuda、prefill_device: cuda。Specifying Running Devices for Modules为所有模块设置兜底设备最后为所有尚未被命中的模块设置兜底设备属性- match: name: ^model\\.layers\\..*\\.|^lm_head replace: class: default kwargs: generate_device: cuda - match: name: ^model.embed_tokens replace: class: default kwargs: generate_device: cpu通过这两条规则把所有未被覆盖的层及其子模块与lm_head放到 cuda把 embedding 放在 cpu。需要注意模块的属性由最先匹配到它的规则决定。例如某注入模块内部已经在replace.kwargs.generate_device中设置了设备那么这条兜底规则中更早设置的设备会优先生效后面的新设置不会覆盖。因此设备相关的 kwargs 应集中规划、避免重复配置造成冲突。从gen_optimize_config的源码逻辑看任何未命中规则的模块都会被写入默认配置{class: default, kwargs: {generate_device: default_device, prefill_device: default_device}}optimize.py默认设备由optimize_and_load_gguf(..., default_devicecuda:0)传入——这正是兜底机制的实现基础。Muti-GPU把 60 层模型拆分到多张显卡若机器有多张 GPU可以按层把不同模块分配到不同显卡。以 DeepSeekV2-Chat60 层为例若有两张卡可将 0-29 层分给cuda:0、30-59 层分给cuda:1。第一步注入模型级算子KDeepseekV2Model并通过transfer_map声明层间切分点。从第 30 层开始激活张量需要跨卡传输transfer_map告诉模型算子在第 30 层把数据搬到cuda:1- match: name: ^model$ replace: class: ktransformers.operators.models.KDeepseekV2Model kwargs: transfer_map: 30: cuda:1官方 multi-gpu 规则中该模块同时带per_layer_prefill_intput_threshold: 00 表示关闭逐层 prefill 流水。第二步为每个模块按层范围分别设置设备。以路由专家为例单卡时规则为- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # Custom MoE kernel with expert parallelism kwargs: generate_device: cuda:0 generate_op: MLPCUDAExperts out_device: cuda:0 recursive: False # Dont recursively inject submodules of this module双卡时需要按层范围拆成两条规则分别把专家计算留在 CPUKExpertsCPU但把输出设备指向各自对应的 GPU# allcate 0-29 layerss out_device to cuda:0 - match: name: ^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: generate_device: cpu generate_op: KExpertsCPU out_device: cuda:0 recursive: False # dont recursively inject submodules of this module # allocate 30-59 layerss out_device to cuda:1 - match: name: ^model\\.layers\\.([345][0-9])\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: generate_device: cpu generate_op: KExpertsCPU out_device: cuda:1 recursive: False # dont recursively inject submodules of this module其中(0|[1-9]|[12][0-9])精确匹配 0-29 的层号([345][0-9])精确匹配 30-59 的层号。其余模块RoPE、线性层、MLA、MoE、兜底 default、lm_head按同样思路以层号正则分组设置generate_device/prefill_device为cuda:0或cuda:1完整示例见仓库中的 DeepSeek-V2-Chat-multi-gpu.yaml含model.norm落cuda:1、lm_head落cuda:1等细节。同理还有 4 卡版本 DeepSeek-V2-Chat-multi-gpu-4.yaml。How to Write a New Operator and Inject into the Model从零编写可注入算子本节以新线性算子为例说明如何编写一个能被注入框架识别并加载权重的算子。第一步继承 BaseInjectedModule所有可注入算子都必须继承BaseInjectedModule定义于 archive/ktransformers/operators/base_operator.py它提供了注入框架所需的属性key、gguf_loader、config、orig_module、prefill_device、generate_device、device并实现了属性访问代理——__getattr__/__setattr__会把未定义属性转发给被替换的原模块orig_module从而保持对外接口兼容。其__init__需要满足如下基本格式class LinearTorchInject(BaseInjectedModule): def __init__( self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module None, generate_device: str cuda, **kwargs, ): super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)第二步自定义参数通过 kwargs 透传如果算子需要额外参数直接写进__init__签名并在 YAML 的replace.kwargs中传入即可。例如新增my_paramclass LinearTorchInject(BaseInjectedModule): def __init__( self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module None, generate_device: str cuda, my_param: bool True, **kwargs, ): super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs) self.my_param my_param对应注入规则- match: name: ^model\\.layers\\..*$ # Regular expression matches the module name. class: torch.nn.Linear # Type restrictions can be added. replace: class: ktransformers.operators.linear.LinearTorchInject # Inject module path kwargs: # Extra parameters generate_device: cuda my_param: True结合 optimize.py 的注入流程可知inject会按key、gguf_loader、config、orig_module位置参数 kwargs关键字参数实例化替换类并把 YAML 中的 kwargs 同时登记进gguf_loader.tensor_device_map用于后续权重加载时确定目标设备。第三步继承 KLinearBase 读取 GGUF 权重线性算子还需要从 GGUF 文件中读取权重。KLinearBase基类archive/ktransformers/operators/linear.py封装了从 GGUF/SafeTensor 加载权重的逻辑其中load_weight方法会根据 loader 类型与张量存在性自动返回nn.Parameter或(weight, bias)元组。用户只需继承并实现load、unload、forward三个函数一个完整可注入的线性类如下class LinearTorchInject(BaseInjectedModule, KLinearBase): def __init__( self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module None, generate_device: str cuda, **kwargs, ): super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs) KLinearBase.__init__(self) self.has_bias False self.dtype torch.get_default_dtype() self.w None self.has_bias False def load(self, w: dict | nn.Parameter | tuple | None None, device: str|None None): if device is None: device self.device if w is None: w self.load_weight(devicedevice) if isinstance(w, nn.Parameter): self.w w.to(dtypeself.dtype).view(self.out_features, self.in_features).T self.has_bias False elif isinstance(w, tuple): self.w w[0].to(dtypeself.dtype).view(self.out_features, self.in_features).T self.bias w[1].to(dtypeself.dtype) self.has_bias True else: raise ValueError(Invalid weight type) self.w self.w.to(device) if self.has_bias: self.bias self.bias.to(device) def unload(self): if self.w is not None: self.w None if self.has_bias: self.bias None def forward(self, x: torch.Tensor) - torch.Tensor: dtype x.dtype out_device x.device x x.to(deviceself.device, dtypeself.dtype) x x self.w if self.has_bias: x x self.bias x x.to(dtypedtype, deviceout_device) return x关键点说明self.load_weight由KLinearBase提供负责把 GGUF 权重加载进模块load中w.to(...).view(self.out_features, self.in_features).T完成权重转置与类型转换GGUF 存储布局与 PyTorch 前向所需布局不同forward中先迁到self.device计算、再迁回out_device与generate_device/out_device的异构设计保持一致unload用于显式释放权重配合del_metaoptimize.py在 meta 设备清理阶段避免残留。BaseInjectedModule.load的默认实现会遍历子模块调用utils.load_weights递归加载权重这也是recursive语义在加载阶段的一种体现。结语从规则到算子的完整注入闭环回顾整个注入流程KTransformers 通过“YAML 规则 →gen_optimize_config生成优化配置 → meta 设备上inject重建模块 →load_weights加载 GGUF 权重 →del_meta清理”五个阶段完成模型改造optimize.py 的optimize_and_load_gguf即完整入口。对于 DeepSeek-V2 家族完整可用的单卡规则与多卡规则分别位于 DeepSeek-V2-Chat.yaml 与 DeepSeek-V2-Chat-multi-gpu.yaml轻量验证可用 DeepSeek-V2-Lite-Chat.yaml。仿照本文的 MLA、专家、线性层、RoPE 与设备指定五类规则并遵循BaseInjectedModule 专用基类的算子编写规范你便可以把自己的自定义算子无缝挂载到任意 transformers 模型的推理链路中。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
KTransformers 操作符注入教程：以 DeepSeek-V2 为例逐层替换自定义算子
KTransformers 操作符注入教程以 DeepSeek-V2 为例逐层替换自定义算子【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers导读KTransformers 的 Inject 注入框架允许用户在不修改 transformers 源码的前提下通过 YAML 规则将模型中的任意子模块替换为自定义算子从而实现 MLA 注意力、量化线性层、CPU 专家并行等异构优化。本教程以 DeepSeekV2-Chat 为例从注入规则语法出发依次演示 MLA 注入、路由专家注入、线性层注入、RoPE 缓冲区补偿、设备指定与多 GPU 划分最后讲解如何基于BaseInjectedModule编写一个可注入的全新算子。读完本文你将能够独立为任意 transformers 模型编写注入规则并把自研算子挂载进 KTransformers 的推理链路。How to Write Injection Rules注入规则的基础语法注入规则本质是一个 YAML 列表KTransformers 的注入框架在模型初始化阶段逐条读取规则对模型树上的每个模块做正则/类型匹配命中后以替换类重建该模块。规则基本形式如下- match: name: ^model\\.layers\\..*\\.*$ # Target module name class: torch.nn.Linear # Target module replace: class: default kwargs: generate_device: cuda:0 # your_op_param_1: 1234 # your_op_param_2: 5678 recursive: True各字段含义match匹配规则包含name与class两种形式。两者可以同时出现也可以单独出现同时出现时要求模块名称与类型均命中才会注入。name正则表达式匹配模块在模型中的完整路径名如model.layers.0.self_attnclass目标模块的 Python 类例如torch.nn.Linear、ktransformers.models.modeling_deepseek.DeepseekV2MoE。replaceclass用于替换目标模块的 Python 类路径。若不想替换仅设置设备等属性设为defaultkwargs替换模块初始化所需的参数字典其中generate_device指定该模块运行的设备可取值cpu、cuda、cuda:1等。recursive是否递归处理该模块的子模块默认True。recursive字段非常关键像 Self-attention 这类模块内部通常包含 q/k/v/o 四个线性子模块如果替换了整个自注意力模块却不想让内部线性层再被其他规则覆盖就需要把该规则设为False典型场景见下文 Routed Experts 注入。从实现看optimize.py 中的gen_optimize_config会深度优先遍历模型树对每个模块依次尝试规则列表class匹配通过isinstance判断name匹配通过re.search判断首个命中的规则生效并break随后inject同文件 L28-L54动态import替换类用module_cls(key..., gguf_loader..., config..., orig_modulechild, **kwargs)构造新模块并通过set_module挂回原位置。因此规则列表中规则的书写顺序即匹配优先级先写的规则优先生效。Understanding Model Structure从权重文件反推模型结构KTransformers 注入框架自由度很高允许你替换/实验任何基础算子但前提是必须清楚所运行模型的内部结构。以deepseek-ai/DeepSeek-V2-Lite-Chat为例打开模型仓库文件列表从.safetensors权重文件名可以看到每一层权重的命名对应注入规则中的match.name从modeling_deepseek.py文件可以看到每个模块类的具体实现对应注入规则中的match.class。由此可得到 DeepSeekV2 的模型骨架model.embed_tokens词嵌入→ 60 层model.layers.N每层含self_attn与mlpmlp内部是gate/experts→model.norm→lm_head。上述self_attn、mlp.experts、lm_head等正是注入规则中正则要命中的路径。KTransformers 官方支持的算子与后端对照表如下对应仓库内ktransformers/operators下的实现matchreplacebackendsdescriptionsLinearKTransformersLinearKLinearMarlinMarlin as backendKLinearTorchpytorch as backendKLinearCPUInferllamafile as backendKLinearFP8Triton fp8_gemm kernel. Requires GPU be able to calculate fp8 dataexpertsKTransformersExpertsKExpertsTorchpytorch as backendKExpertsMarlinMarlin as backendKExpertsCPUllamafile as backendAttentionKDeepseekV2AttentionKDeepseekV2AttentionMLA implementationMoEKMistralSparseMoEBlockKQwen2MoeSparseMoeBlockMoE for Qwen2KDeepseekV2MoEKDeepseekV2MoEMoE for DeepseekV2ModelKQwen2MoeModelKQwen2MoeModelModel for Qwen2KDeepseekV2ModelKDeepseekV2ModelModel for DeepseekV2RoPERotaryEmbeddingRotaryEmbeddingRoPE moduleYarnRotaryEmbeddingYarnRotaryEmbeddingRoPE module在 archive/ktransformers/operators/attention.py 中可以看到KDeepseekV2Attention以(BaseInjectedModule, DeepseekV2Attention)双重继承的方式实现experts.py 中KTransformersExperts通过EXPERTS_MAP[generate_op]分发到KExpertsCPU/KExpertsTorch等具体后端linear.py 中KTransformersLinear同样以generate_op选择KLinearMarlin等实现。generate_op就是你在 YAML 里指定具体后端的关键字。下面以 DeepSeekV2-Chat 为目标逐步注入自定义 Marlin 线性层、基于 Absorption 的 MLA 注意力、自定义专家模块、自定义 MoE 模块、自定义 RoPE 模块并为每个模块指定运行设备。Matrix Absorption-based MLA Injection替换注意力为 MLA 实现MLAMulti-head Latent Attention通过矩阵吸收把 q/k/v 投影合并显著降低 KV Cache 占用。注入 Attention 模块只需用正则命中 transformers 中的模块名替换为 KTransformers 的 MLA 实现- match: name: ^model\\.layers\\..*\\.self_attn$ # Regular expression replace: class: ktransformers.operators.attention.KDeepseekV2Attention # Optimized MLA implementation规则中match负责指定被替换模块replace指定注入的模块类及初始化关键字。完整的DeepSeek-V2-Chat.yaml中还为该规则补充了设备参数- match: name: ^model\\.layers\\..*\\.self_attn$ replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: cuda prefill_device: cudaprefill_device与generate_device分别控制预填充prefill阶段与逐 token 生成generate阶段的运行设备两阶段可以异构这是 KTransformers 异构推理的核心设计之一见 base_operator.py 中BaseInjectedModule对这两个属性的登记。Injection of Routed Experts注入路由专家并改写 MoE 前向对于图中标注为exps的路由专家Routed Experts注入的是被包装模块KTransformersExperts包裹的 CPUInfer 算子。KTransformersExperts有多个后端实现需要通过关键字告诉包装模块选择哪个实现以及如何使用它。transformers 中 MoE 用nn.ModuleList实现专家列表我们不希望 KTransformers 遍历列表逐个注入子模块因此在规则中设置recursive: False- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # Custom MoE kernel with expert parallelism kwargs: generate_device: cpu generate_op: MLPCPUExperts out_device: cuda recursive: False # Dont recursively inject submodules of this modulegenerate_device: cpu专家权重与计算放在 CPU由 llamafile/CPUInfer 承载generate_op: MLPCPUExperts选择 CPU 专家后端out_device: cuda专家输出写回 GPU供后续算子使用。实际仓库中的DeepSeek-V2-Chat.yaml同时为 prefill 阶段配置了prefill_op: KExpertsTorch、prefill_device: cuda即 prefill 用 GPU 上的 PyTorch 专家、generate 用 CPU 专家实现两阶段后端分离- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism 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注意官方 YAML 中 generate 阶段用的后端名是KExpertsCPU即文档正文中的MLPCPUExperts演变而来对应 experts.py 中out_device参数“输出应落在哪个设备”的语义——代码中KExpertsCPU.output_gpu_map[out_device]会为每个目标 GPU 预分配输出缓冲并通过submit_with_cuda_stream与 CUDA 流异步协同。注入路由专家为自定义模块后原nn.ModuleList的接口不再可用因此必须修改 FFN 模块的forward函数。最简单的做法是实现一个带自定义forward的新模块并注入- match: class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # MLP module with custom forward functionKDeepseekV2MoE定义在 archive/ktransformers/operators/experts.py它重写了 MoE 的前向逻辑以适配被替换后的专家模块调用方式。Injection of Linear Layers用类型检查精确注入量化线性层对于剩余线性层我们希望使用量化算子来节省存储并提升性能。由于目前没有 MLA 与量化联合使用的研究我们不希望把线性层注入到 MLA 算子内部因此通过正则负向断言 类型检查双重约束只有 name 与 class同时命中的模块才会被注入- match: name: ^model\\.layers\\.(?!.*self_attn).*$ # 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 kwargs: generate_device: cuda generate_op: QuantizedLinearMarlin(?!.*self_attn)是正则负向前瞻排除所有包含self_attn的路径从而保证 MLA 模块内部不被线性层规则波及。官方DeepSeek-V2-Chat.yaml中该规则更精确地排除了kv_b_projMLA 吸收合并后不再单独需要它并分别指定了生成与预填充阶段的后端- match: name: ^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$ # 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 kwargs: generate_device: cuda prefill_device: cuda generate_op: KLinearMarlin prefill_op: KLinearTorch即生成阶段使用 Marlin 量化内核KLinearMarlin预填充阶段退回 PyTorch 实现KLinearTorch。KLinearMarlin封装了 GPTQ 量化权重与 Marlin 反量化内核linear.py 的KLinearBase基类会在初始化时从orig_module或 GGUF 权重元信息中解析in_features/out_features并维护loaded状态。Injection of Modules with Pre-calculated Buffers补偿 RoPE 的预计算缓冲区为避免初始化注入后的原模型占用资源KTransformers 使用torch.device(meta)元设备初始化原模型见 optimize.pywith torch.device(meta): inject(...)随后再load_weights并del_meta清理 meta 参数。但 RoPE 模块在初始化时会预计算一批缓冲区如频率表元设备下不会真正执行计算因此需要在加载模型时补偿这些缓冲区的计算。做法是向旋转嵌入模块注入自定义模块在加载阶段完成预计算- match: class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbeddingktransformers.operators.RoPE.YarnRotaryEmbedding定义于 archive/ktransformers/operators/RoPE.py继承了原DeepseekV2YarnRotaryEmbedding的语义但在实际设备上完成缓冲区预计算从而让基于 YARN 的长上下文旋转位置编码在注入框架下正常工作。官方规则同时为其指定了generate_device: cuda、prefill_device: cuda。Specifying Running Devices for Modules为所有模块设置兜底设备最后为所有尚未被命中的模块设置兜底设备属性- match: name: ^model\\.layers\\..*\\.|^lm_head replace: class: default kwargs: generate_device: cuda - match: name: ^model.embed_tokens replace: class: default kwargs: generate_device: cpu通过这两条规则把所有未被覆盖的层及其子模块与lm_head放到 cuda把 embedding 放在 cpu。需要注意模块的属性由最先匹配到它的规则决定。例如某注入模块内部已经在replace.kwargs.generate_device中设置了设备那么这条兜底规则中更早设置的设备会优先生效后面的新设置不会覆盖。因此设备相关的 kwargs 应集中规划、避免重复配置造成冲突。从gen_optimize_config的源码逻辑看任何未命中规则的模块都会被写入默认配置{class: default, kwargs: {generate_device: default_device, prefill_device: default_device}}optimize.py默认设备由optimize_and_load_gguf(..., default_devicecuda:0)传入——这正是兜底机制的实现基础。Muti-GPU把 60 层模型拆分到多张显卡若机器有多张 GPU可以按层把不同模块分配到不同显卡。以 DeepSeekV2-Chat60 层为例若有两张卡可将 0-29 层分给cuda:0、30-59 层分给cuda:1。第一步注入模型级算子KDeepseekV2Model并通过transfer_map声明层间切分点。从第 30 层开始激活张量需要跨卡传输transfer_map告诉模型算子在第 30 层把数据搬到cuda:1- match: name: ^model$ replace: class: ktransformers.operators.models.KDeepseekV2Model kwargs: transfer_map: 30: cuda:1官方 multi-gpu 规则中该模块同时带per_layer_prefill_intput_threshold: 00 表示关闭逐层 prefill 流水。第二步为每个模块按层范围分别设置设备。以路由专家为例单卡时规则为- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # Custom MoE kernel with expert parallelism kwargs: generate_device: cuda:0 generate_op: MLPCUDAExperts out_device: cuda:0 recursive: False # Dont recursively inject submodules of this module双卡时需要按层范围拆成两条规则分别把专家计算留在 CPUKExpertsCPU但把输出设备指向各自对应的 GPU# allcate 0-29 layerss out_device to cuda:0 - match: name: ^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: generate_device: cpu generate_op: KExpertsCPU out_device: cuda:0 recursive: False # dont recursively inject submodules of this module # allocate 30-59 layerss out_device to cuda:1 - match: name: ^model\\.layers\\.([345][0-9])\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: generate_device: cpu generate_op: KExpertsCPU out_device: cuda:1 recursive: False # dont recursively inject submodules of this module其中(0|[1-9]|[12][0-9])精确匹配 0-29 的层号([345][0-9])精确匹配 30-59 的层号。其余模块RoPE、线性层、MLA、MoE、兜底 default、lm_head按同样思路以层号正则分组设置generate_device/prefill_device为cuda:0或cuda:1完整示例见仓库中的 DeepSeek-V2-Chat-multi-gpu.yaml含model.norm落cuda:1、lm_head落cuda:1等细节。同理还有 4 卡版本 DeepSeek-V2-Chat-multi-gpu-4.yaml。How to Write a New Operator and Inject into the Model从零编写可注入算子本节以新线性算子为例说明如何编写一个能被注入框架识别并加载权重的算子。第一步继承 BaseInjectedModule所有可注入算子都必须继承BaseInjectedModule定义于 archive/ktransformers/operators/base_operator.py它提供了注入框架所需的属性key、gguf_loader、config、orig_module、prefill_device、generate_device、device并实现了属性访问代理——__getattr__/__setattr__会把未定义属性转发给被替换的原模块orig_module从而保持对外接口兼容。其__init__需要满足如下基本格式class LinearTorchInject(BaseInjectedModule): def __init__( self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module None, generate_device: str cuda, **kwargs, ): super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs)第二步自定义参数通过 kwargs 透传如果算子需要额外参数直接写进__init__签名并在 YAML 的replace.kwargs中传入即可。例如新增my_paramclass LinearTorchInject(BaseInjectedModule): def __init__( self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module None, generate_device: str cuda, my_param: bool True, **kwargs, ): super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs) self.my_param my_param对应注入规则- match: name: ^model\\.layers\\..*$ # Regular expression matches the module name. class: torch.nn.Linear # Type restrictions can be added. replace: class: ktransformers.operators.linear.LinearTorchInject # Inject module path kwargs: # Extra parameters generate_device: cuda my_param: True结合 optimize.py 的注入流程可知inject会按key、gguf_loader、config、orig_module位置参数 kwargs关键字参数实例化替换类并把 YAML 中的 kwargs 同时登记进gguf_loader.tensor_device_map用于后续权重加载时确定目标设备。第三步继承 KLinearBase 读取 GGUF 权重线性算子还需要从 GGUF 文件中读取权重。KLinearBase基类archive/ktransformers/operators/linear.py封装了从 GGUF/SafeTensor 加载权重的逻辑其中load_weight方法会根据 loader 类型与张量存在性自动返回nn.Parameter或(weight, bias)元组。用户只需继承并实现load、unload、forward三个函数一个完整可注入的线性类如下class LinearTorchInject(BaseInjectedModule, KLinearBase): def __init__( self, key: str, gguf_loader: GGUFLoader, config: PretrainedConfig, orig_module: nn.Module None, generate_device: str cuda, **kwargs, ): super().__init__(key, gguf_loader, config, orig_module, generate_device, **kwargs) KLinearBase.__init__(self) self.has_bias False self.dtype torch.get_default_dtype() self.w None self.has_bias False def load(self, w: dict | nn.Parameter | tuple | None None, device: str|None None): if device is None: device self.device if w is None: w self.load_weight(devicedevice) if isinstance(w, nn.Parameter): self.w w.to(dtypeself.dtype).view(self.out_features, self.in_features).T self.has_bias False elif isinstance(w, tuple): self.w w[0].to(dtypeself.dtype).view(self.out_features, self.in_features).T self.bias w[1].to(dtypeself.dtype) self.has_bias True else: raise ValueError(Invalid weight type) self.w self.w.to(device) if self.has_bias: self.bias self.bias.to(device) def unload(self): if self.w is not None: self.w None if self.has_bias: self.bias None def forward(self, x: torch.Tensor) - torch.Tensor: dtype x.dtype out_device x.device x x.to(deviceself.device, dtypeself.dtype) x x self.w if self.has_bias: x x self.bias x x.to(dtypedtype, deviceout_device) return x关键点说明self.load_weight由KLinearBase提供负责把 GGUF 权重加载进模块load中w.to(...).view(self.out_features, self.in_features).T完成权重转置与类型转换GGUF 存储布局与 PyTorch 前向所需布局不同forward中先迁到self.device计算、再迁回out_device与generate_device/out_device的异构设计保持一致unload用于显式释放权重配合del_metaoptimize.py在 meta 设备清理阶段避免残留。BaseInjectedModule.load的默认实现会遍历子模块调用utils.load_weights递归加载权重这也是recursive语义在加载阶段的一种体现。结语从规则到算子的完整注入闭环回顾整个注入流程KTransformers 通过“YAML 规则 →gen_optimize_config生成优化配置 → meta 设备上inject重建模块 →load_weights加载 GGUF 权重 →del_meta清理”五个阶段完成模型改造optimize.py 的optimize_and_load_gguf即完整入口。对于 DeepSeek-V2 家族完整可用的单卡规则与多卡规则分别位于 DeepSeek-V2-Chat.yaml 与 DeepSeek-V2-Chat-multi-gpu.yaml轻量验证可用 DeepSeek-V2-Lite-Chat.yaml。仿照本文的 MLA、专家、线性层、RoPE 与设备指定五类规则并遵循BaseInjectedModule 专用基类的算子编写规范你便可以把自己的自定义算子无缝挂载到任意 transformers 模型的推理链路中。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考