MoE的VRAM磁盘缓存使单块DGX Spark上的Kimi 2.7达到340 pp/s和9.6 tg/s

该策略有效利用VRAM作为磁盘缓存，使MoE专家留在llama.cpp的CUDA计算路径上。先看数据。详细解释见下文。 DGX Spark上的数据：Kimi-K2.7-Code.i1-IQ_S.gguf 204GB 1T.A32B https://huggingface.co/mradermacher/Kimi-K2.7-Code-i1-GGUF 运行方法 pp512 tg128 备注 A cuda, -ot regex A 见下文, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH 340.02 ± 37.75 9.58 ± 0.07 所有张量在统一内存上，专家通过内存映射保持在可分页主机内存中，最后几层固定在RAM中 B cuda, -ot regex B 见下文, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH 154.31 ± 30.62 8.75 ± 0.22 A，但最后几层未固定 C cuda, -ot regex C 见下文 GGML_CUDA_ENABLE_UNIFIED_MEMORY 222.20 ± 66.04 3.26 ± 0.01 B，但移除最小批次标志 D cuda, GGML_CUDA_ENABLE_UNIFIED_MEMORY 崩溃 崩溃 C，但移除专家卸载 E cpu mmap 4.23 ± 0.22 1.63 ± 0.66 未涉及CUDA -ot regex A: '^(?!blk.(5[5-9]|60).ffn(down|gate|up)_exps.weight$).*.ffn(down|gate|up)_exps.weight=CPU' -ot regex B: '.*\.ffn_(down|gate|up)_exps\.weight=CPU' -ot regex C: '.*.ffn_(down|gate|up)_exps.weight=CPU' 3090s PCIe 4.0 + 128GB DDR4上的数据：Minimax-M2.7-K_G_3.00.gguf 80GB 230B.A10B https://huggingface.co/Goldkoron/MiniMax-M2.7 运行方法 pp512 tg128 备注 A cuda, -ot regex A 见下文, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH 54.03 ± 3.85 2.90 ± 0.09 与DGX上B相同的策略 B cuda -ngl 17 73.80 ± 2.10 1.66 ± 0.02 普通层拆分 C 3090x2, cuda, -ot regex C 见下文, GGML_CUDA_ENABLE_UNIFIED_MEMORY, GGML_OP_OFFLOAD_MIN_BATCH 48.86 ± 0.53 2.40 ± 0.01 A，但使用两块3090 -ot regex A: '^(?!blk.(5[5-9]|60).ffn(down|gate|up)_exps.weight$).*.ffn(down|gate|up)_exps.weight=CPU' -ot regex C: '^(?!blk.(5[5-9]|60).ffn(down|gate|up)_exps.weight$).*.ffn(down|gate|up)_exps.weight=CPU' 背景 我想在单块DGX Spark上运行超过128GB的模型。MoE模型提供了绝佳的机会，因为如果激活参数能完全用CUDA计算，则运行时无需将所有模型权重存储在GPU RAM中。经过实验，我找到了一组llama.cpp的设置来快速运行Kimi 2.7。 设置 llama-bench的通用设置：-r 50, -fa on, -mmp 1, 关闭系统交换。 llama.cpp b10075 编译时启用了kleidai支持，获胜策略如下： GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 GGML_OP_OFFLOAD_MIN_BATCH=1 ./llama-bench -m ~/Kimi-K2.7-Code.i1-IQ_S.gguf -fa on -mmp 1 -ot '^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU' -r 50. 我使用 cmake -B build -DGGMLCUDA=ON -DGGML_CPU_KLEIDIAI=ON 配置了 b10075，并用 cmake --build build --config Release 编译。 工作原理 首先，如果 GGML_CUDA_ENABLE_UNIFIED_MEMORY 为 1，CUDA设备缓冲区使用 cudaMallocManaged，每个张量只有一个虚拟地址，对GPU和CPU都有效。当CUDA内核访问已存在于VRAM中的页面时，会正常执行；如果不在VRAM中，CPU RAM中的页面会将页面迁移到VRAM；如果既不在GPU RAM也不在CPU RAM中，则通过 mmap 从磁盘将页面获取到CPU RAM，然后发送到VRAM。在DGX Spark上，由于GPU RAM和CPU RAM是统一的，cudaMallocManaged的开销要低得多。 其次，对于 GGML_OP_OFFLOAD_MIN_BATCH（默认值为32），如果批次大小小于阈值且权重在CPU RAM上，则在CPU上执行，否则在GPU上执行。如果将其设置为1，则强制所有计算在GPU上进行。这有助于令牌生成，因为在tg中不会达到该阈值。 第三，当指定 -ot '.*.ffn_(down|gate|up)_exps.weight=CPU' 时，专家权重将被放在CPU RAM上，并可能因 mmap 而被换出。这使得冷专家权重更可能留在磁盘上。 最后，'^(?!blk\.(5[5-9]|60)\.ffn_(down|gate|up)_exps\.weight$).*\.ffn_(down|gate|up)_exps\.weight=CPU' 用于将最后几个块保留在GPU RAM上，因为这些块具有更多的动态路由，当RAM空间足够大时，将它们保留在GPU RAM上可以更快地避免开销。 统一RAM场景 对于DGX Spark，它有较大的内存池，因此我可以将某些专家块固定在VRAM中（运行A）。如果不固定，速度会变慢（运行B）。如果移除 GGML_OP_OFFLOAD_MIN_BATCH 标志，令牌生成会变慢（运行C），因为专家前向传播在CPU上执行。如果进一步移除 GGML_CUDA_ENABLE_UNIFIED_MEMORY 标志（运行D），会发生OOM，因为GPU RAM无法通过 mmap 机制换出。最后，如果所有计算都在CPU上进行，速度最慢（运行E）。 在运行B与运行C之间出现了一个有趣的权衡：GGML_OP_OFFLOAD_MIN_BATCH 显著提高了tg，但略微影响了pp。可能需要进行代码修改以缓解这一权衡。 非统一RAM硬件场景 对于3090 + DDR4，我尚未深入测试，但该策略也基本适用。我没有将最后几个块放在VRAM上，因为24GB太小。可以看到运行A的tg高于普通 ngl 拆分（运行B），因为专家在GPU上计算。但提示处理效果较差，可能是由于PCIe开销。我还测试了两块3090 + DDR4，但总体上更慢，可能是PCIe传输成本太高。我认为适当调整 -ot 卸载可能有机会在双GPU场景下改善。 注意 Mac Apple Silicon很诱人，因为它具有快速统一RAM，但我找不到简单的方法禁用交换。所提出的方法可能因RAM交换而导致大量SSD写入，可能不是好主意，因为它可能增加SSD磨损。 在新模型上执行此操作的建议方法是先卸载专家（DGX Spark中的运行B策略），然后逐步将更多专家固定到VRAM。 我期待看到即将推出的Kimi K3量化是否可以同样应用。相同的策略可跨统一内存和非统一内存架构移植。3090+DDR4设置下的tg也有所改进。 编辑：修正格式