KTransformers DeepSeek-V2 异构 MoE 推理从 MLA 优化算子到 YAML 注入规则实战【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers本文基于 KTransformers 仓库中的 DeepSeek-V2 注入教程完整讲清三件事为什么 236B 参数的 DeepSeek-V2 能跑在 21GB 显存的桌面显卡上KTransformers 为此实现了哪些底层算子优化版 MLA、Marlin/llamafile 量化内核以及如何用一套 YAML 注入规则把这些算子精准替换进 Transformers 模型。读完本文你能独立编写和修改注入规则文件理解match/replace/recursive语义与设备分派参数的底层含义并为自定义算子接入 KTransformers 注入框架。背景236B MoE 模型本地部署的资源困境DeepSeek-(Code)-V2 是一系列强大的混合专家MoE模型总参数 236B每个 token 激活约 21B 参数。与采用 GQA 的早期模型不同DeepSeek-V2 引入了多头潜在注意力Multi-head Latent Attention, MLA它用一个共享的、联合压缩的表示来表征所有 KV 头从而显著减小推理时的 KV cache 体积。然而实际部署并不友好官方文档表明标准推理需要 8 张 80GB GPU即便 Q4_k_m 量化版本也至少需要 2 张 80GB GPU。KTransformers 的目标就是打破这一限制——通过一组组合优化在仅有21GB VRAM 136GB DRAM的桌面机器上运行完整模型。KTransformers 采用的优化可归纳为三类优化版 MLA 算子不展开压缩表示直接在压缩空间做注意力高级量化内核GPU 侧用 MarlinCPU 侧用基于 llamafile 的 CPUInfer直接对量化数据类型计算算术强度引导的算子卸载按算术强度把算子分配到 GPU/CPU。下面逐一展开并给出仓库中对应的源码实现。优化一吸收分解矩阵的 MLA 算子DeepSeek-V2 的 MLA 算子把 KV 头表示为一个共享的压缩表示compressed KV但官方开源实现会把压缩表示显式解压后再缓存解压后的 KV 对。这一过程既放大了 KV cache 体积又降低了推理性能。KTransformers 的做法来自 MLA 原论文把解压矩阵直接吸收到q_proj与out_proj的权重中。这样注意力计算全程不需要解压压缩表示KV cache 大幅缩小算子算术强度显著提升从而更高效地利用 GPU 张量核。仓库源码中这一优化落在 KDeepseekV2Attention。核心是get_absorbed()方法它把kv_b_proj的权重按头切分并 reshape分别得到q_absorb用于把 query 投到压缩空间与out_absorb用于把压缩空间的注意力输出还原回 v_head_dim# archive/ktransformers/operators/attention.pyL69-L75 def get_absorbed(self) - Tuple[torch.Tensor, torch.Tensor]: if not (hasattr(self, q_absorb) and hasattr(self, out_absorb)): kv_b_proj self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank) self.q_absorb kv_b_proj[:, :self.qk_nope_head_dim, :].view(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank) self.out_absorb kv_b_proj[:, self.qk_nope_head_dim:, :].view(self.num_heads, self.v_head_dim, self.kv_lora_rank) return self.q_absorb, self.out_absorb在forward_linux_triton路径attention.py#L196中decode 阶段只把 query 的 nope 部分乘上q_absorb使其进入kv_lora_rank维的压缩空间然后与缓存中的compressed_kv、k_pe直接计算注意力输出再乘out_absorb还原。KV cache 中存放的始终是压缩表示compressed_kvk_pe配合仓库自研的分页缓存custom_cache.py与 Triton 分页注意力内核triton_attention.py完成高效解码。forward()attention.py#L685-L743会根据运行环境自动分派到 flashinfer、Triton 或分块实现。优化二面向量化数据类型的计算内核DeepSeek-V2 原始权重为 BF16约 470GB 存储超出主流台式机内存。KTransformers 直接复用 GGUF 社区的量化权重如 Q4_K_M简化用户流程。难点在于主流高性能 BLAS 不支持量化数据类型原始 Torch 实现必须先把张量反量化回支持类型才能计算带来额外算力与内存带宽开销。KTransformers 的方案是直接对量化数据类型做计算GPU 侧Marlin。专为现代 GPU 架构设计可充分发挥 tensor core 能力。按 Marlin 论文的基准数据其相对对应 Torch 实现可获得接近理想的 3.87 倍加速原文档引用。CPU 侧llamafileCPUInfer。利用现代 CPU 指令扩展如 AVX512-BF16AMD Zen4 及更新与 AVX-VNNIIntel Alder Lake 及更新并且内置专家并行等 MoE 推理优化。KTransformers 的微基准显示低比特表示下 CPUInfer 推理比 Torch 快数倍而且由于 Torch 基线要么用 BF16/FP16 权重占更多内存要么在反量化时变慢实际差距更大。在源码中这些内核统一封装在 linear.py 里KTransformersLinearL906是注入模块内部按 prefill/generate 两种模式选择KLinearTorchL158、KLinearMarlinL595或KLinearCPUInferL723。专家部分则由KExpertsCPUexperts.py#L143基于 llamafile 实现。优化三算术强度引导的算子卸载把 236B 参数全部放进显存显然不现实。KTransformers 的原则是按算子算术强度排序把最密集的算子尽量放在 GPU。MLA 算子优化后 MLA 含 128 个共享压缩 KV 表示的头算术强度达512是模型中最密集的算子小 batch 推理时尤为突出因此分配给 GPU 以利用 tensor core。路由专家routed experts每个 Transformer block 含 160 个 MoE 专家占总参数 96%但 MoE router 每个 token 只激活 6 个即解码阶段只用到 MoE 参数的 3.75%。batch size 为 1 时MoE 运算的算术强度约0.075本质上是一组批量 GEMV矩阵-向量乘非常适合 CPU 处理。按此原则KTransformers 把 MoE 参数与词嵌入放到 CPU利用其大容量内存把共享专家、注意力模块的投影层与 MLA 留在 GPU 显存——这些参数被每个 token 访问放在 GPU 能最大化高带宽优势。最终 Q4_K_M 版本的资源占用约为20.7GB VRAM 136GB DRAM普通桌面即可运行且可按实际配置在同一原则下调整布局。YAML 注入框架优化如何落到模型上要落地上述优化用户需要编写一个包含优化规则的 YAML 文件。KTransformers 会遍历模型所有子模块按规则匹配并替换为高级模块。注入引擎的底层实现注入流程由 optimize.py 驱动入口是optimize_and_load_gguf()L129-L163整体分四步规则匹配gen_optimize_config()L67-L118在 meta 设备上递归遍历模型对每个模块依次检查规则。match部分支持name对模块全路径做正则搜索re.search与classisinstance检查两种条件两者同时给出时必须同时满足第一个命中的规则即生效并停止后续匹配递归控制规则可用recursive: False阻止继续深入该模块的子模块对nn.ModuleList形式的专家列表至关重要见下文模块替换inject()L28-L54动态导入替换类构造时传入key模块路径、gguf_loader、config、orig_module被替换的原模块以及规则里的kwargs然后set_module完成替换权重加载通过GGUFLoader从 GGUF 文件按tensor_device_map把权重加载到各注入模块指定设备上最后del_meta()清理 meta 设备上的占位参数并释放缓存。规则一注入优化 MLA 注意力只需按 Transformers 中的模块名做正则匹配替换为预实现的优化算子- match: name: ^model\\.layers\\..*\\.self_attn$ # 正则表达式 replace: class: ktransformers.operators.attention.KDeepseekV2Attention # 优化的 MLA 实现每条规则由match与replace两部分组成match指明被替换的模块replace指明注入的模块及初始化参数。规则二注入路由专家CPUInfer 专家并行注入的专家模块是 CPUInfer 的封装类KTransformersExpertsexperts.py#L686-L747。封装内含多种实现需要用参数指定prefill 阶段用什么、generate 阶段用什么、输出落在哪个设备。从源码可以看到封装的内部机制__init__按prefill_op/generate_op分别实例化两套专家后端prefill_experts、generate_expertsload()/set_inference_mode()在 prefill 与 generate 两种推理模式间切换时执行加载/卸载卸载 prefill 后端、加载 generate 后端等forward()根据当前模式分派到对应后端。值得注意的是prefill 与 generate 的差异化配置仅在启用 layer-wise prefill逐层 prefill时生效否则 prefill 与 generate 使用相同配置。- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # 带专家并行的自定义 MoE 内核 device: cpu # 初始化时加载到的设备 kwargs: prefill_device: cuda prefill_op: KExpertsTorch generate_device: cpu generate_op: KExpertsCPU out_device: cuda recursive: False # 不再递归注入该模块的子模块其中recursive: False是刻意设计Transformers 中 MoE 专家用nn.ModuleList实现若不禁止递归注入引擎会遍历列表中每个专家子模块导致注入失败或产生冗余替换。把整个专家列表换成自定义模块后就不能再用nn.ModuleList的默认接口需要修改 FFN 模块的 forward。KTransformers 的做法是实现一个新的 MLP 模块并注入。此时可以用class而非name匹配待替换模块- match: class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # 带自定义 forward 的 MLP 模块对应的KDeepseekV2MoEexperts.py#L874自定义了 forward先过gate得到 top-k 索引与权重单 token 且处于 CUDA graph 捕获时走submit_for_one_decode异步路径提交 CPU 推理任务并同步取回结果否则按 prefill/generate 模式分派。规则三其余 Linear 模块换成 Marlin 量化内核对剩余 Linear 模块希望统一用量化内核但不希望替换 MLA 算子内部的 Linear当时尚不清楚量化对 MLA 效果的影响。因此规则同时给出正则排除self_attn与类检查torch.nn.Linear只有 name 与 class 同时命中的模块才被替换- match: name: ^model\\.layers\\.(?!.*self_attn).*$ # 正则排除 self_attn 下的模块 class: torch.nn.Linear # 仅同时满足 name 与 class 的模块会被注入 replace: class: ktransformers.operators.linear.KTransformersLinear # 量化数据类型上的优化内核 kwargs: generate_device: cuda prefill_device: cuda generate_op: KLinearMarlin prefill_op: KLinearTorch仓库中现成的参考配置 DeepSeek-V2-Chat.yaml 在此基础上更完整可视为可直接使用的模板。它在上面几条规则之外还包含RoPE 预计算缓冲区替换见下一条Linear 替换的正则细化为排除self_attn.kv_b_proj而非整个self_attn并为lm_head增加独立规则model注入KDeepseekV2Model并用per_layer_prefill_intput_threshold: 0关闭 layer-wise prefill0 表示关闭embed_tokens用class: default标记——不替换模块仅把该词嵌入指定到cpu上与词嵌入放 CPU的卸载策略一致。规则四RoPE 预计算缓冲区原模型在 meta 设备上初始化而 rotary embedding 模块在初始化时预计算缓冲区——在 meta 设备上这是空操作什么都不算。因此必须在加载模型时真正计算这些缓冲区。KTransformers 注入自定义 RoPE 模块在加载时执行预计算- match: class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding对应实现见 RoPE.py。包装你的自定义模块KTransformers 已实现了一批算子但框架是开放的你同样可以注入自己写的模块。要做的只有两件事——包装自定义模块 编写 YAML 规则。框架提供了基础算子 BaseInjectedModule自定义模块继承它后按需重写__init__、forward或load__init__保存注入与运行所需的信息模块 key、GGUF 加载器、config、原始模块、prefill/generate 设备。重写时必须在自己的初始化里调用基类的__init__。从源码看基类还用object.__setattr__管理这些字段并通过__getattr__/__setattr__把未定义的属性代理到orig_module使注入模块对外暴露与原模块一致的接口base_operator.py#L31-L56forward推理时被调用的入口作者可自由实现以获得更高性能load负责加载该模块的全部参数。默认实现是递归调用所有子模块的load可以重写它来自定义加载方式、显式控制各子模块的加载设备、内核选择等。加载与运行 DeepSeek-V2配置好 YAML 规则与 GGUF 量化权重后推理入口在 local_chat.py模型先在 meta 设备构建随后调用optimize_and_load_gguf(model, optimize_config_path, gguf_path, config)local_chat.py#L139完成规则匹配 → 模块注入 → 权重加载全流程再根据gguf_loader.tensor_device_map初始化缓存与推理流程。运行前提具备 CUDA GPU教程场景为约 21GB 显存 136GB 内存的桌面机、已下载对应的 DeepSeek-V2 GGUF 量化权重Q4_K_M 为教程基准。需要说明的是本文对应的注入实现与规则文件在当前仓库中位于archive/ktransformers/目录下如 optimize.py、operators/ 与 optimize_rules/YAML 中的类导入路径以ktransformers.为包名前缀按该版本的包结构解析。小结DeepSeek-V2 能在 21GB VRAM 136GB DRAM 的桌面上运行依赖三项优化的组合吸收分解矩阵的 MLA 算子KV cache 更小、算术强度 512、Marlin/llamafile 量化直算内核GPU 3.87x 级别加速、CPU 数倍加速、以及按算术强度MoE 仅 0.075把 96% 参数卸载到 CPUYAML 注入框架把换哪个模块、换成什么、用什么参数变成声明式规则match支持 name 正则与 class 双重条件、首个命中规则生效、recursive控制递归深度replace通过classkwargsprefill_device/generate_device/prefill_op/generate_op/out_device完成算子与设备分派框架以BaseInjectedModule的__init__/forward/load三个接口为扩展点为替换官方算子、注入自研内核提供了标准路径。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 KTransformers DeepSeek-V2 异构 MoE 推理：从 MLA 优化算子到 YAML 注入规则实战
KTransformers DeepSeek-V2 异构 MoE 推理从 MLA 优化算子到 YAML 注入规则实战【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers本文基于 KTransformers 仓库中的 DeepSeek-V2 注入教程完整讲清三件事为什么 236B 参数的 DeepSeek-V2 能跑在 21GB 显存的桌面显卡上KTransformers 为此实现了哪些底层算子优化版 MLA、Marlin/llamafile 量化内核以及如何用一套 YAML 注入规则把这些算子精准替换进 Transformers 模型。读完本文你能独立编写和修改注入规则文件理解match/replace/recursive语义与设备分派参数的底层含义并为自定义算子接入 KTransformers 注入框架。背景236B MoE 模型本地部署的资源困境DeepSeek-(Code)-V2 是一系列强大的混合专家MoE模型总参数 236B每个 token 激活约 21B 参数。与采用 GQA 的早期模型不同DeepSeek-V2 引入了多头潜在注意力Multi-head Latent Attention, MLA它用一个共享的、联合压缩的表示来表征所有 KV 头从而显著减小推理时的 KV cache 体积。然而实际部署并不友好官方文档表明标准推理需要 8 张 80GB GPU即便 Q4_k_m 量化版本也至少需要 2 张 80GB GPU。KTransformers 的目标就是打破这一限制——通过一组组合优化在仅有21GB VRAM 136GB DRAM的桌面机器上运行完整模型。KTransformers 采用的优化可归纳为三类优化版 MLA 算子不展开压缩表示直接在压缩空间做注意力高级量化内核GPU 侧用 MarlinCPU 侧用基于 llamafile 的 CPUInfer直接对量化数据类型计算算术强度引导的算子卸载按算术强度把算子分配到 GPU/CPU。下面逐一展开并给出仓库中对应的源码实现。优化一吸收分解矩阵的 MLA 算子DeepSeek-V2 的 MLA 算子把 KV 头表示为一个共享的压缩表示compressed KV但官方开源实现会把压缩表示显式解压后再缓存解压后的 KV 对。这一过程既放大了 KV cache 体积又降低了推理性能。KTransformers 的做法来自 MLA 原论文把解压矩阵直接吸收到q_proj与out_proj的权重中。这样注意力计算全程不需要解压压缩表示KV cache 大幅缩小算子算术强度显著提升从而更高效地利用 GPU 张量核。仓库源码中这一优化落在 KDeepseekV2Attention。核心是get_absorbed()方法它把kv_b_proj的权重按头切分并 reshape分别得到q_absorb用于把 query 投到压缩空间与out_absorb用于把压缩空间的注意力输出还原回 v_head_dim# archive/ktransformers/operators/attention.pyL69-L75 def get_absorbed(self) - Tuple[torch.Tensor, torch.Tensor]: if not (hasattr(self, q_absorb) and hasattr(self, out_absorb)): kv_b_proj self.kv_b_proj.weight.view(self.num_heads, -1, self.kv_lora_rank) self.q_absorb kv_b_proj[:, :self.qk_nope_head_dim, :].view(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank) self.out_absorb kv_b_proj[:, self.qk_nope_head_dim:, :].view(self.num_heads, self.v_head_dim, self.kv_lora_rank) return self.q_absorb, self.out_absorb在forward_linux_triton路径attention.py#L196中decode 阶段只把 query 的 nope 部分乘上q_absorb使其进入kv_lora_rank维的压缩空间然后与缓存中的compressed_kv、k_pe直接计算注意力输出再乘out_absorb还原。KV cache 中存放的始终是压缩表示compressed_kvk_pe配合仓库自研的分页缓存custom_cache.py与 Triton 分页注意力内核triton_attention.py完成高效解码。forward()attention.py#L685-L743会根据运行环境自动分派到 flashinfer、Triton 或分块实现。优化二面向量化数据类型的计算内核DeepSeek-V2 原始权重为 BF16约 470GB 存储超出主流台式机内存。KTransformers 直接复用 GGUF 社区的量化权重如 Q4_K_M简化用户流程。难点在于主流高性能 BLAS 不支持量化数据类型原始 Torch 实现必须先把张量反量化回支持类型才能计算带来额外算力与内存带宽开销。KTransformers 的方案是直接对量化数据类型做计算GPU 侧Marlin。专为现代 GPU 架构设计可充分发挥 tensor core 能力。按 Marlin 论文的基准数据其相对对应 Torch 实现可获得接近理想的 3.87 倍加速原文档引用。CPU 侧llamafileCPUInfer。利用现代 CPU 指令扩展如 AVX512-BF16AMD Zen4 及更新与 AVX-VNNIIntel Alder Lake 及更新并且内置专家并行等 MoE 推理优化。KTransformers 的微基准显示低比特表示下 CPUInfer 推理比 Torch 快数倍而且由于 Torch 基线要么用 BF16/FP16 权重占更多内存要么在反量化时变慢实际差距更大。在源码中这些内核统一封装在 linear.py 里KTransformersLinearL906是注入模块内部按 prefill/generate 两种模式选择KLinearTorchL158、KLinearMarlinL595或KLinearCPUInferL723。专家部分则由KExpertsCPUexperts.py#L143基于 llamafile 实现。优化三算术强度引导的算子卸载把 236B 参数全部放进显存显然不现实。KTransformers 的原则是按算子算术强度排序把最密集的算子尽量放在 GPU。MLA 算子优化后 MLA 含 128 个共享压缩 KV 表示的头算术强度达512是模型中最密集的算子小 batch 推理时尤为突出因此分配给 GPU 以利用 tensor core。路由专家routed experts每个 Transformer block 含 160 个 MoE 专家占总参数 96%但 MoE router 每个 token 只激活 6 个即解码阶段只用到 MoE 参数的 3.75%。batch size 为 1 时MoE 运算的算术强度约0.075本质上是一组批量 GEMV矩阵-向量乘非常适合 CPU 处理。按此原则KTransformers 把 MoE 参数与词嵌入放到 CPU利用其大容量内存把共享专家、注意力模块的投影层与 MLA 留在 GPU 显存——这些参数被每个 token 访问放在 GPU 能最大化高带宽优势。最终 Q4_K_M 版本的资源占用约为20.7GB VRAM 136GB DRAM普通桌面即可运行且可按实际配置在同一原则下调整布局。YAML 注入框架优化如何落到模型上要落地上述优化用户需要编写一个包含优化规则的 YAML 文件。KTransformers 会遍历模型所有子模块按规则匹配并替换为高级模块。注入引擎的底层实现注入流程由 optimize.py 驱动入口是optimize_and_load_gguf()L129-L163整体分四步规则匹配gen_optimize_config()L67-L118在 meta 设备上递归遍历模型对每个模块依次检查规则。match部分支持name对模块全路径做正则搜索re.search与classisinstance检查两种条件两者同时给出时必须同时满足第一个命中的规则即生效并停止后续匹配递归控制规则可用recursive: False阻止继续深入该模块的子模块对nn.ModuleList形式的专家列表至关重要见下文模块替换inject()L28-L54动态导入替换类构造时传入key模块路径、gguf_loader、config、orig_module被替换的原模块以及规则里的kwargs然后set_module完成替换权重加载通过GGUFLoader从 GGUF 文件按tensor_device_map把权重加载到各注入模块指定设备上最后del_meta()清理 meta 设备上的占位参数并释放缓存。规则一注入优化 MLA 注意力只需按 Transformers 中的模块名做正则匹配替换为预实现的优化算子- match: name: ^model\\.layers\\..*\\.self_attn$ # 正则表达式 replace: class: ktransformers.operators.attention.KDeepseekV2Attention # 优化的 MLA 实现每条规则由match与replace两部分组成match指明被替换的模块replace指明注入的模块及初始化参数。规则二注入路由专家CPUInfer 专家并行注入的专家模块是 CPUInfer 的封装类KTransformersExpertsexperts.py#L686-L747。封装内含多种实现需要用参数指定prefill 阶段用什么、generate 阶段用什么、输出落在哪个设备。从源码可以看到封装的内部机制__init__按prefill_op/generate_op分别实例化两套专家后端prefill_experts、generate_expertsload()/set_inference_mode()在 prefill 与 generate 两种推理模式间切换时执行加载/卸载卸载 prefill 后端、加载 generate 后端等forward()根据当前模式分派到对应后端。值得注意的是prefill 与 generate 的差异化配置仅在启用 layer-wise prefill逐层 prefill时生效否则 prefill 与 generate 使用相同配置。- match: name: ^model\\.layers\\..*\\.mlp\\.experts$ replace: class: ktransformers.operators.experts.KTransformersExperts # 带专家并行的自定义 MoE 内核 device: cpu # 初始化时加载到的设备 kwargs: prefill_device: cuda prefill_op: KExpertsTorch generate_device: cpu generate_op: KExpertsCPU out_device: cuda recursive: False # 不再递归注入该模块的子模块其中recursive: False是刻意设计Transformers 中 MoE 专家用nn.ModuleList实现若不禁止递归注入引擎会遍历列表中每个专家子模块导致注入失败或产生冗余替换。把整个专家列表换成自定义模块后就不能再用nn.ModuleList的默认接口需要修改 FFN 模块的 forward。KTransformers 的做法是实现一个新的 MLP 模块并注入。此时可以用class而非name匹配待替换模块- match: class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # 带自定义 forward 的 MLP 模块对应的KDeepseekV2MoEexperts.py#L874自定义了 forward先过gate得到 top-k 索引与权重单 token 且处于 CUDA graph 捕获时走submit_for_one_decode异步路径提交 CPU 推理任务并同步取回结果否则按 prefill/generate 模式分派。规则三其余 Linear 模块换成 Marlin 量化内核对剩余 Linear 模块希望统一用量化内核但不希望替换 MLA 算子内部的 Linear当时尚不清楚量化对 MLA 效果的影响。因此规则同时给出正则排除self_attn与类检查torch.nn.Linear只有 name 与 class 同时命中的模块才被替换- match: name: ^model\\.layers\\.(?!.*self_attn).*$ # 正则排除 self_attn 下的模块 class: torch.nn.Linear # 仅同时满足 name 与 class 的模块会被注入 replace: class: ktransformers.operators.linear.KTransformersLinear # 量化数据类型上的优化内核 kwargs: generate_device: cuda prefill_device: cuda generate_op: KLinearMarlin prefill_op: KLinearTorch仓库中现成的参考配置 DeepSeek-V2-Chat.yaml 在此基础上更完整可视为可直接使用的模板。它在上面几条规则之外还包含RoPE 预计算缓冲区替换见下一条Linear 替换的正则细化为排除self_attn.kv_b_proj而非整个self_attn并为lm_head增加独立规则model注入KDeepseekV2Model并用per_layer_prefill_intput_threshold: 0关闭 layer-wise prefill0 表示关闭embed_tokens用class: default标记——不替换模块仅把该词嵌入指定到cpu上与词嵌入放 CPU的卸载策略一致。规则四RoPE 预计算缓冲区原模型在 meta 设备上初始化而 rotary embedding 模块在初始化时预计算缓冲区——在 meta 设备上这是空操作什么都不算。因此必须在加载模型时真正计算这些缓冲区。KTransformers 注入自定义 RoPE 模块在加载时执行预计算- match: class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding对应实现见 RoPE.py。包装你的自定义模块KTransformers 已实现了一批算子但框架是开放的你同样可以注入自己写的模块。要做的只有两件事——包装自定义模块 编写 YAML 规则。框架提供了基础算子 BaseInjectedModule自定义模块继承它后按需重写__init__、forward或load__init__保存注入与运行所需的信息模块 key、GGUF 加载器、config、原始模块、prefill/generate 设备。重写时必须在自己的初始化里调用基类的__init__。从源码看基类还用object.__setattr__管理这些字段并通过__getattr__/__setattr__把未定义的属性代理到orig_module使注入模块对外暴露与原模块一致的接口base_operator.py#L31-L56forward推理时被调用的入口作者可自由实现以获得更高性能load负责加载该模块的全部参数。默认实现是递归调用所有子模块的load可以重写它来自定义加载方式、显式控制各子模块的加载设备、内核选择等。加载与运行 DeepSeek-V2配置好 YAML 规则与 GGUF 量化权重后推理入口在 local_chat.py模型先在 meta 设备构建随后调用optimize_and_load_gguf(model, optimize_config_path, gguf_path, config)local_chat.py#L139完成规则匹配 → 模块注入 → 权重加载全流程再根据gguf_loader.tensor_device_map初始化缓存与推理流程。运行前提具备 CUDA GPU教程场景为约 21GB 显存 136GB 内存的桌面机、已下载对应的 DeepSeek-V2 GGUF 量化权重Q4_K_M 为教程基准。需要说明的是本文对应的注入实现与规则文件在当前仓库中位于archive/ktransformers/目录下如 optimize.py、operators/ 与 optimize_rules/YAML 中的类导入路径以ktransformers.为包名前缀按该版本的包结构解析。小结DeepSeek-V2 能在 21GB VRAM 136GB DRAM 的桌面上运行依赖三项优化的组合吸收分解矩阵的 MLA 算子KV cache 更小、算术强度 512、Marlin/llamafile 量化直算内核GPU 3.87x 级别加速、CPU 数倍加速、以及按算术强度MoE 仅 0.075把 96% 参数卸载到 CPUYAML 注入框架把换哪个模块、换成什么、用什么参数变成声明式规则match支持 name 正则与 class 双重条件、首个命中规则生效、recursive控制递归深度replace通过classkwargsprefill_device/generate_device/prefill_op/generate_op/out_device完成算子与设备分派框架以BaseInjectedModule的__init__/forward/load三个接口为扩展点为替换官方算子、注入自研内核提供了标准路径。【免费下载链接】ktransformersA Flexible Framework for Experiencing Heterogeneous LLM Inference/Fine-tune Optimizations项目地址: https://gitcode.com/GitHub_Trending/ktr/ktransformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考