背景

MoE 模型的总参数量虽然很大, 但每个 token 实际只会激活其中少量 experts, 因此 active parameters 远小于总参数量, 这个现象称为 sparse activation; 它让消费级 GPU 具备了计算大规模 MoE 模型的可能性, 因为 GPU 只需要执行当前激活的部分, 不需要同时计算完整的 expert pool;

但是, 完整 expert pool 仍然可能远大于 GPU 显存, 因此需要主要保存在 host memory 中; FreeToken 希望把个人电脑的 GPU, CPU, memory 和 interconnect 组合成一个完整的 serving system; 当需要的 expert 没有命中 GPU cache 时, 系统可以将 expert weights 搬到 GPU 上计算, 也可以将较小的 activation tensor 传到 CPU 侧直接计算;

FreeToken 希望感知 CPU, GPU 和 memory 之间的实际 bandwidth, 并据此动态选择执行方式, 从而尽可能利用个人电脑已有的异构资源, 提高大规模 MoE 模型在个人 PC 上的推理性能;

现有系统的痛点

在 prefill 阶段, 一个长 prompt 中的不同 tokens 会被路由到不同 experts; 虽然单个 token 只激活少量 experts, 但所有 tokens 的路由取并集后, 每一层往往会覆盖接近完整的 expert set, 使原本稀疏的 working set 接近 dense; 由于完整 expert pool 无法常驻 GPU, prefill 需要通过 PCIe 从 host memory 传输几乎整层的 expert weights; 如果系统按需读取这些 experts, GPU 就会在等待数据时闲置;

个人电脑的资源也不会被 serving system 独占; 桌面系统, 浏览器和游戏都可能随时占用 GPU compute 与 VRAM, 因而系统实际可用的 memory budget 会动态变化; 随着 agent session 的 context 不断增长, KV cache 需要的空间也会增加, 所以 VRAM 在 KV cache 和 expert cache 之间的最佳划分同样不是固定的; 一个只在启动时决定资源分配的静态策略, 很难适应个人电脑实际的 workflow;