KTransformers 是一个灵活的、以 Python 为中心的框架,旨在通过先进的内核优化和放置 / 并行策略提升 Hugging Face Transformers 的使用体验。它具有高度的可扩展性,用户可通过单行代码注入优化模块,获得兼容 Transformers 的接口、符合 OpenAI 和 Ollama 的 RESTful API,甚至简化的 ChatGPT 风格的 Web UI。
KTransformers的性能优化基本囊括了目前主流的优化手段,包括:
本地运行 6710 亿参数的 DeepSeek - Coder - V3/R1:使用仅 14GB VRAM 和 382GB DRAM 运行其 Q4_K_M 版本。
****说明一下:不支持并发,项目实现了单卡体验的可能性,但绝对上不了生产!太慢!用4个自行车轱辘加上2节干电池做了一个4轮电动车,但是,但是,但是,绝对不等于==智能电动汽车!!!!
预填充(Prefill)速度 (tokens/s) :
KTransformers: 54.21 (32 核心) → 74.362 (双插槽,2×32 核心) → 255.26 (优化的 AMX 基 MoE 内核,仅 V0.3) → 286.55 (选择性使用 6 个专家,仅 V0.3)
与 llama.cpp 在 2×32 核心下 10.31 tokens/s 相比,速度提升高达 27.79 倍
解码(Decode)速度 (tokens/s):
KTransformers: 8.73 (32 核心) → 11.26 (双插槽, 2×32 核心) → 13.69 (选择性使用 6 个专家,仅 V0.3)
与 llama.cpp 在 2×32 核心下 4.51 tokens/s 相比,速度提升高达 3.03 倍
最佳性能测试(V0.2):
CPU:Intel ® Xeon ® Gold 6454S 1T 内存 (2 NUMA 节点)
GPU:4090D 24G 显存
内存:标准 DDR5 - 4800 服务器内存 (1 TB)
“6 个专家” 情况是 V0.3 预览版中内容
Prompt (500 tokens) |
双插槽 Ktrans (6 个专家) | 双插槽 Ktrans (8 个专家) | Single socket Ktrans (6 个专家) | Single socket Ktrans (8 个专家) | llama.cpp (8 个专家) |
---|---|---|---|---|---|
预填充(Prefill) token/s | 97.32 | 82.94 | 65.14 | 54.21 | 10.31 |
解码(Decode) token/s | 13.69 | 12.208 | 10.303 | 8.73 | 4.51 |
最高加速比在解码方面达到 3.03x 倍,在预填充方面达到 9.44x 倍。
Prompt length | 1K | 2K | 4K | 8K |
---|---|---|---|---|
KTrans (8 个专家) Prefill token/s | 185.96 | 255.26 | 252.58 | 195.62 |
KTrans (6 个专家) Prefill token/s | 203.70 | 286.55 | 271.08 | 207.20 |
KTrans V0.3 的预填充速度比 KTrans V0.2 快 3.45x 倍,比 llama.cpp 快 27.79x 倍。
解码速度与 KTrans V0.2(6 个专家版本)相同,因此省略。
主要加速来自于
从我们对 DeepSeekV2、DeepSeekV3 和 DeepSeekR1 的研究中,当我们略微减少推理中的激活专家数量时,输出质量没有变化。但解码和预填充的速度加快了,这令人鼓舞。因此,我们的展示利用了这一发现。
我们的 local_chat 测试命令是:
git clone https://github.com/kvcache-ai/ktransformers.git
cd ktransformers
git submodule init
git submodule update
numactl -N 1 -m 1 python ./ktransformers/local_chat.py --model_path <your model path> --gguf_path <your gguf path> --prompt_file <your prompt txt file> --cpu_infer 33 --max_new_tokens 1000
<当您看到聊天时,按回车键加载文本提示文件>
可以是本地路径,也可以是在线路径,例如 deepseek-ai/DeepSeek-V3。如果在线连接出现问题,可以尝试使用镜像(hf-mirror.com)
也可以是在线路径,但由于其体积较大,我们建议您下载并量化模型(注意这是目录路径)
--max_new_tokens 1000
是最大输出 token 长度。如果发现答案被截断,可以增加此数字以获得更长的答案(但要注意内存不足问题,增加此数字会降低生成速度).
命令 numactl -N 1 -m 1 的目的是避免 NUMA 节点之间的数据传输
注意!如果测试 R1 可能会跳过思考。因此,可以添加参数:--force_think true
,这在 常见问题解答 部分中解释。
在安装之前(使用 install.sh 或 make dev_install
),请确保设置环境变量 USE_NUMA=1
,方法是 export USE_NUMA=1
(如果已经安装,请重新安装并设置此环境变量)
我们的 local_chat 测试命令是:
git clone https://github.com/kvcache-ai/ktransformers.git
cd ktransformers
git submodule init
git submodule update
export USE_NUMA=1
make dev_install # or sh ./install.sh
python ./ktransformers/local_chat.py --model_path <your model path> --gguf_path <your gguf path> --prompt_file <your prompt txt file> --cpu_infer 65 --max_new_tokens 1000
<当您看到聊天时,按回车键加载文本提示文件>
参数的含义相同。但因为我们使用双插槽,所以将 cpu_infer 设置为 65。
我们的 local_chat 测试命令是:
wget https://github.com/kvcache-ai/ktransformers/releases/download/v0.1.4/ktransformers-0.3.0rc0+cu126torch26fancy-cp311-cp311-linux_x86_64.whl
pip install ./ktransformers-0.3.0rc0+cu126torch26fancy-cp311-cp311-linux_x86_64.whl
python -m ktransformers.local_chat --model_path <your model path> --gguf_path <your gguf path> --prompt_file <your prompt txt file> --cpu_infer 65 --max_new_tokens 1000
<当您看到聊天时,按回车键加载文本提示文件>
参数的含义与 V0.2 相同。但因为我们使用双插槽,所以将 cpu_infer 设置为 65。
命令参数 --cpu_infer 65
指定使用多少核心(超过物理核心数量是可以的,但并不是越多越好。根据实际核心数量适当降低此值)。
为什么使用 CPU/GPU 混合推理?
DeepSeek 的 MLA 操作符计算密集。虽然全部在 CPU 上运行是可行的,但将繁重的计算任务卸载到 GPU 上能带来巨大的性能提升。
加速来自哪里?
为什么选择英特尔 CPU?
英特尔目前是唯一支持 AMX 类似指令的 CPU 供应商,与仅支持 AVX 的替代方案相比,性能显著更好。
KTransformers 与 vLLM 有何不同?
vLLM 是一个用于大规模部署优化的出色框架。与之不同,KTransformers 特别专注于受资源限制的本地部署场景。
KTransformers 着重挖掘异构计算的潜力,比如在量化模型的 GPU/CPU 卸载方面进行优化。
举例来说,它分别支持用于 CPU 和 GPU 的高效 Llamafile 和 Marlin 内核,以此提升本地推理性能
model_path指向的是原始的模型路径,而gguf_path指向的是量化后的GGUF格式模型文件。 为什么要这么做?
KTransformers的优化策略,包括将部分参数卸载到CPU并使用GGUF格式的量化权重,而GPU部分则使用Marlin内核处理。因此,model_path可能用于加载模型的结构和部分参数,而gguf_path提供量化后的权重,用于CPU端的计算。所以,KTransformers部署时需要同时指定model_path和gguf_path:
如果没有足够的VRAM,但我有多个GPU,该如何利用它们?多GPU会提升推理速度吗?
修改 YAML 注入模板
KTransformers 使用 YAML 注入模板来优化和配置模型。对于多 GPU 设置,你需要在 YAML 模板中指定设备分配。以下是一个示例,展示如何将模块分配到不同的 GPU 上:
- match:
name: "^model\\.layers\\.0$" # 匹配模型的第一层
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformerLinear
device: "cuda:0" # 将此模块分配到GPU 0
kwargs:
generate_device: "cuda:0"
generate_linear_type: "QuantizedLinearMarlin"
- match:
name: "^model\\.layers\\.1$" # 匹配模型的第二层
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformerLinear
device: "cuda:1" # 将此模块分配到GPU 1
kwargs:
generate_device: "cuda:1"
generate_linear_type: "QuantizedLinearMarlin"
在这个示例中,模型的不同层被分配到了不同的 GPU 上。你可以根据模型的结构和你的 GPU 数量,扩展这个模板以分配更多的模块。
确保使用optimize_and_load_gguf
函数,并传入正确的 YAML 配置文件
optimize_and_load_gguf
函数会根据 YAML 模板中的配置,将模型模块分配到相应的 GPU 上。 注意: ktransformers的多GPU策略为pipline,无法加速模型的推理,仅用于模型的权重分配。GPU/CPU混合推理deepseekR1使用了多GPU,但推理性能和单GPU一致,没有提升,多卡并不会加速推理。https://github.com/kvcache-ai/ktransformers/issues/345,用8卡4090跑的,速度还没有单卡的快
如何获得最佳性能?
必须设置–cpu_infer为要使用的核数。使用的核越多,模型运行速度越快,但并非越多越好。将其调整得略低一些,以适应实际核数。例如,–cpu_infer 65指定使用多少个核,超过物理数量是可以的,但也不是越多越好,稍微调低到实际的核心数量。
如果我获得的VRAM(也就是显存)比模型要求的多,我该如何充分利用它?
加大上下文,max_new_tokenslocal_chat.py可以以通过设置更大的值来增加上下文窗口的大小。
修改 YAML 注入模板,KTransformers 使用 YAML 注入模板来优化和配置模型。对于多 GPU 设置,你需要在 YAML 模板中指定设备分配。以下是一个示例,展示如何将模块分配到不同的 GPU 上:
- match:
name: "^model\\.layers\\.0$" # 匹配模型的第一层
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformerLinear
device: "cuda:0" # 将此模块分配到GPU 01,
kwargs:
generate_device: "cuda:0"
generate_linear_type: "QuantizedLinearMarlin"
- match:
name: "^model\\.layers\\.1$" # 匹配模型的第二层
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformerLinear
device: "cuda:1" # 将此模块分配到GPU 1
kwargs:
generate_device: "cuda:1"
generate_linear_type: "QuantizedLinearMarlin"
在这个示例中,模型的不同层被分配到了不同的 GPU 上。你可以根据模型的结构和你的 GPU 数量,扩展这个模板以分配更多的模块。
KTransformers依赖Python环境,推荐使用Python 3.8及以上版本。你可以从Python官方网站下载并安装Python。
install.sh
脚本。该脚本会自动安装项目所需的依赖项,包括torch
、transformers
等。install.bat
脚本。它会执行类似的操作,安装Windows环境下所需的依赖包。pip
命令。首先创建并激活虚拟环境(可选但推荐):# 创建虚拟环境(假设使用venv)
python -m venv myenv
# 激活虚拟环境
# 在Linux/Mac上
source myenv/bin/activate
# 在Windows上
myenv\Scripts\activate
然后安装依赖项:
pip install -r requirements.txt
requirements.txt
文件包含了KTransformers运行所需的所有Python包及其版本要求。
安装完依赖项后,可以通过以下方式安装KTransformers:
pip install ktransformers
git clone https://github.com/kvcache-ai/ktransformers.git
然后进入项目目录并安装:
cd ktransformers
pip install -e.
-e
选项表示以可编辑模式安装,这样在你修改源代码时,无需重新安装即可生效,方便开发和调试。
安装完成后,可以通过运行简单的示例脚本来验证KTransformers是否正确安装。在项目的示例目录中,有一些示例代码,例如example.py
。运行该脚本:
python example.py
如果一切正常,脚本应该能够成功导入KTransformers并执行相应的操作,比如加载模型并进行推理。若遇到问题,请参考常见问题解答(FAQ)或在GitHub仓库中提交问题以获取帮助。
KTransformers支持多GPU配置,这可以显著提升大型语言模型(LLM)的推理速度。本教程将指导你如何在KTransformers中设置和使用多GPU进行推理。
在使用多GPU之前,你可以通过运行以下Python代码片段来检查系统中可用的GPU:
import torch
if torch.cuda.is_available():
num_gpus = torch.cuda.device_count()
print(f"Detected {num_gpus} GPUs.")
else:
print("No GPUs detected. Please check your GPU installation.")
如果代码正确检测到多个GPU,你就可以继续进行下一步配置。
KTransformers使用YAML注入模板来优化和配置模型。对于多GPU设置,你需要在YAML模板中指定设备分配。以下是一个示例,展示如何将模块分配到不同的GPU上:
- match:
name: "^model\\.layers\\.0$" # 匹配模型的第一层
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformerLinear
device: "cuda:0" # 将此模块分配到GPU 0
kwargs:
generate_device: "cuda:0"
generate_linear_type: "QuantizedLinearMarlin"
- match:
name: "^model\\.layers\\.1$" # 匹配模型的第二层
class: torch.nn.Linear
replace:
class: ktransformers.operators.linear.KTransformerLinear
device: "cuda:1" # 将此模块分配到GPU 1
kwargs:
generate_device: "cuda:1"
generate_linear_type: "QuantizedLinearMarlin"
在这个示例中,模型的不同层被分配到了不同的GPU上。你可以根据模型的结构和你的GPU数量,扩展这个模板以分配更多的模块。
在代码中加载模型时,确保使用optimize_and_load_gguf
函数,并传入正确的YAML配置文件。以下是一个示例:
import torch
from transformers import AutoModelForCausalLM
from ktransformers import optimize_and_load_gguf
with torch.device("meta"):
model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
optimize_and_load_gguf(model, optimize_rule_path="path/to/your/yaml/file.yaml", gguf_path="path/to/your/model.gguf", config=config)
optimize_and_load_gguf
函数会根据YAML模板中的配置,将模型模块分配到相应的GPU上。
加载模型后,你可以像往常一样进行推理。KTransformers提供了prefill_and_generate
函数,它针对多GPU设置进行了优化,可以进一步提高推理速度:
from ktransformers import prefill_and_generate
import torch
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("your_model_tokenizer")
input_text = "Your input prompt here"
input_tensor = tokenizer.encode(input_text, return_tensors="pt").cuda()
generated = prefill_and_generate(model, tokenizer, input_tensor, max_new_tokens=1000)
print(generated)
这个函数会在多个GPU上并行处理输入,生成输出文本。