kvcache-ai/ktransformers
1
0
Fork 0
mirror of https://github.com/kvcache-ai/ktransformers.git synced 2026-05-11 08:20:21 +00:00
Code Issues Packages Projects Releases Wiki Activity

update marlin expert example

Browse Source
This commit is contained in:
Azure
2025-02-12 04:11:00 +00:00
parent a2fc2a8658
commit 0564ac8465
1 changed files with 20 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
ktransformers/optimize/optimize_rules/DeepSeek-V3-Chat-multi-gpu-marlin.yaml
View File

@@ -79,6 +79,24 @@
generate_device: "cuda:1"
prefill_device: "cuda:1"
- match:
name: "^model\\.layers\\.(0|[1-4])\\.mlp\\.experts$" # inject experts in layer 0~4 as marlin expert
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
generate_device: "cuda:0" # run in cuda:0
generate_op: "KExpertsMarlin"
recursive: False
- match:
name: "^model\\.layers\\.([3][0])\\.mlp\\.experts$" # inject experts in layer 30~31 as marlin expert
replace:
class: ktransformers.operators.experts.KTransformersExperts
kwargs:
generate_device: "cuda:1"
generate_op: "KExpertsMarlin"
recursive: False
- match:
name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$"
replace:
@@ -139,5 +157,5 @@
replace:
class: "default"
kwargs:
generate_device: "cuda:1"
prefill_device: "cuda:1"
generate_device: "cuda:0"
prefill_device: "cuda:0"