ds4.c 学习笔记

Appearance

SSD 流式加载——让超出 RAM 容量的模型通过"热专家缓存 + 冷专家按需从 SSD 加载"的方式运行，把"能不能跑"从硬截止变成速度的连续谱。

为什么需要 SSD Streaming ​

DeepSeek V4 Flash 的 IQ2_XXS 模型约 81GB，128GB MacBook 可以完整加载。但 PRO Q4 模型需要 ~838GB，远超单机内存。即使 Flash 模型，如果用户只有 64GB RAM，也无法运行。

SSD Streaming 的核心洞察：MoE 层有 256 个路由专家，但每个 token 只激活 6 个。路由专家占模型大部分空间，但大部分时间都在睡觉。把"热"专家放 RAM，"冷"专家留磁盘按需加载，就能在有限 RAM 下运行超大模型。

传统: RAM 必须 ≥ 模型大小 → 硬截止，差 1GB 也不行
SSD:  RAM 越多 → 缓存越多 → 越快，RAM 少 → 缓存少 → 慢但能跑

核心原理 ​

仍需 RAM 的部分：共享专家、KV cache、embedding、Router、output head——这些每个 token 都要用。只有路由专家的权重可以留在 SSD 上。

缓存计划（Cache Plan） ​

ds4_ssd_cache_plan 计算给定 RAM 预算下能缓存多少完整专家：

c
typedef struct {
    uint64_t model_target_bytes;    // 推荐字节的 80%
    uint64_t cache_bytes;           // 总预算减去非路由部分
    uint64_t effective_cache_bytes; // 实际使用的缓存字节
    uint32_t cache_experts;         // 可缓存的专家数量
} ds4_ssd_cache_plan;

// 自动计算：给定推荐工作集 → 扣除非路由权重 → 算能放几个专家
ds4_ssd_auto_cache_plan(recommended_bytes, non_routed_bytes, per_expert_bytes, max_model_experts, &out);

目标：80% 的推荐内存用于缓存，留 20% 给 KV cache 和工作缓冲区。

显式指定的 --ssd-streaming-cache-experts 32GB 不是通用字节缓存——引擎把它换算成"当前 GGUF 能放下多少个完整路由专家"（gate + up + down 一起算）。32GB 对 Flash IQ2/Q2 GGUF 约等于 4854 个路由专家。只有自动预算模式才会替你做"扣除非路由权重"的减法；显式设值时，非路由权重、KV cache、graph scratch 和激活所需的内存需要你自己留余量。

缓存上限与 mlock 余量 ​

显式请求的 NGB 缓存如果太大，会在推理开始前被封顶——否则 macOS 会拒绝 mlock 锁定或系统卡死。封顶的判据是"专家缓冲区必须保持可锁定（lockable）"，而不是"尽量满足用户请求"。

当系统处于额外内存压力下、mlock 仍然失败时，引擎不会硬塞可换页（pageable）的专家缓存条目，而是：

  1. 释放一段已锁定的缓存余量（locked-cache margin）；
  2. 用实测可锁定的缓存大小继续运行。

这条策略比"要么全锁、要么报错退出"更实用：内存压力是动态的，宁可缓存小一点也别让系统陷入换页泥潭。环境变量 DS4_METAL_STREAM_EXPERT_CACHE_*（及等价的 CUDA/ROCm 路径）可以观测这些决策。

启动热度排名表（Hotlist） ​

ds4_streaming_hotlist.inc（13,334 行）是一个预生成的静态热度排名表，按流行度排列所有 (layer, expert_id) 对。启动时预加载最热门的专家，确保常用专家在第一次交互前已在缓存中。

自动预热默认封顶 4096 个专家——避免启动时花太久把缓存填满。需要测量更大预热规模时才用 --ssd-streaming-preload-experts N 显式指定；正常使用应保留默认预热，不要加 --ssd-streaming-cold（冷启动跳过预热，首次交互会很慢）。

运行期路由热度淘汰（Route Hotness） ​

Hotlist 是启动时的静态预热；真正决定运行期淘汰谁的是一张运行期热度表 g_stream_expert_cache_route_hotness[layer][expert]。它取代了纯 LRU：

c
// 每处理 16 个 decode token，整张热度表右移一位（>>= 1）衰减
enum { DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS = 16 };

static void ds4_gpu_stream_expert_cache_decay_route_hotness(void) {
    for (layer ...) for (expert ...)
        g_stream_expert_cache_route_hotness[layer][expert] >>= 1;
}

淘汰时优先驱逐热度最低的专家。衰减粒度 16 token 是关键参数：比"逐 token 衰减"便宜得多（整张表一次性移位），又比"从不衰减"更适应工作负载变化——一串 token 反复命中某专家会让它迅速变热，而冷却下来的专家会被新热点挤出去。

作用域是"提示局部"（prompt-local）：每次新 prompt 开始（start == 0）时调用 ds4_gpu_stream_expert_cache_reset_route_hotness() 清零热度表——因为新对话的路由分布与上一轮无关。但常驻缓存本身有意跨会话保持热。ds4_gpu.h 的注释把这条策略讲得很明白：

c
/* Reset only the prompt-local eviction heuristic.  The resident SSD expert
 * cache itself is intentionally kept warm across sessions. */
void ds4_gpu_stream_expert_cache_reset_route_hotness(void);

为什么：热度表只描述"这一轮里谁该被淘汰"，是会过时的启发式，新 prompt 必须重置；而缓存里已加载的专家是真金白银加载过的数据，跨会话复用零成本，没有理由清掉。区分"易变的启发式状态"和"昂贵的物化缓存"是这条设计的核心。

三路 Prefill 调度 ​

短 prompt 在 streaming 下很慢（每层的固定开销主导），因此默认 prefill 调度器根据 prompt 长度走三条路径：

Prompt 规模策略原因
极短逐 token decode 式 prefill批量加载专家的开销不划算
中等selected-expert 批量 prefill只加载被选中的专家，按需
长全层 layer-major prefill一次处理多层、多 token，摊薄 SSD 读取

切分点按量化格式调优（metal_graph_use_streaming_decode_prefill 的 prefill 截止阈值），避免短 prompt 落到"加载一堆用不上的专家"的慢路径。这也是为什么短 prompt streaming 比 long prefill 慢得多——长 prefill 一层能处理很多 token，缓存命中率高得多。

混合精度路由专家（Mixed-Precision Routed Experts） ​

当 GGUF 的路由专家跨层不统一量化时（例如全局 IQ2_XXS/Q2_K，但有几层用 --tensor-type 提升到 Q4_K，即"per-layer boosted quants"），streaming 路径曾经每个请求都失败：

  • batch-selected-addr prefill 用 decode 静态 span 集映射"提升层"（没有 exps 张量）；
  • Metal encoder 的 IQ2-only gate 把非 IQ2 层丢给 full-tensor wrap_model_range 兜底；
  • 结果没有任何已安装视图覆盖该范围（Metal model range ... is not covered by mapped model views）。

更隐蔽的是：专家缓存的单尺寸 slab 分配器会被"超大专家"毒化——note_expert_size 用 last-writer-wins 记录专家尺寸，预算被扭曲，封顶后可能 slab 复用饥饿。

修复分两部分：

  1. 检测提升层（weights_streaming_layer_experts_uniform：某层的 per-expert 字节数 ≠ 首个路由层推出的 slab 尺寸类），把这些层的 ffn_{gate,up,down}_exps 张量纳入 decode-span 构建器（model_map_span_vec_include_layer_decode），让现有兜底路径能覆盖它们。
  2. 启动时预置 slab 尺寸类（ds4_gpu_set_streaming_expert_cache_expert_bytes），并让 note_expert_size 改为 freeze+reject 而非 last-writer-wins——超大尺寸的层确定性绕过缓存、走 mapped-view 路径，不污染 slab 分配器。

启动日志会打印 uniform/boosted 的分布，并在"多数路由层都偏离尺寸类"时告警。均匀模型按构造与改动前逐字节相同（helper 对每层都返回 true，不加额外 span，预置值等于第一次 note 会写入的值）。

GPU 后端集成 ​

SSD streaming 现已跨三个 GPU 后端实现，统一在 ds4_gpu.h 抽象层下：

  • Metal（ds4_metal.m）：最完整的实现，专家缓存管理、异步加载、路由热度淘汰、三路 prefill。
  • CUDA（ds4_cuda.cu）：SSD streaming 完整实现（不再是 stub），含有界缓存分配和进度上报，与 PRO/Flash 路径集成。
  • ROCm（ds4_rocm.cu + rocm/*.cuh）：selected-expert 缓存、重叠读取、全层 streaming prefill、常驻缓存查找、ROCm 缓存复用处理。

核心 GPU 抽象层函数（streaming API 已重构为传一张 ds4_gpu_stream_expert_table 表，避免长参数列表）：

c
// 重构后：把 model_map/offsets/bytes 打包成一张表，跨后端共用
typedef struct ds4_gpu_stream_expert_table {
    const void *model_map;   uint64_t model_size;
    uint32_t layer;          uint32_t n_total_expert;
    uint64_t gate_offset, up_offset, down_offset;
    uint64_t gate_expert_bytes, down_expert_bytes;
} ds4_gpu_stream_expert_table;

int ds4_gpu_stream_expert_cache_begin_selected_load(
        const ds4_gpu_stream_expert_table *table,
        const int32_t *selected_ids, uint32_t n_selected);  // 异步加载缺失专家
uint32_t ds4_gpu_stream_expert_cache_current_count(void);   // 当前缓存专家数

#ifdef DS4_ROCM_BUILD 守卫了 ROCm 独有的 layer-major 批量加载钩子（ds4_gpu_stream_expert_cache_load_layer 等），非 ROCm 后端提供 no-op stub 保持链接通过。

内存模拟锁（Memory Lock） ​

ds4_ssd_memory_lock 用 mmap + mlock 模拟已占用内存，用于测试：

c
// --simulate-used-memory 32GB 模拟系统已用 32GB
// 分 256 MiB 块逐块 mlock，避免长时间不可中断
ds4_ssd_memory_lock_acquire(lock);

命令行选项 ​

选项说明
--ssd-streaming启用 SSD streaming 模式
--ssd-streaming-cache-experts N指定缓存专家数量（32GB 等字节预算，非专家个数；过大时推理前自动封顶）
--ssd-streaming-preload-experts N显式指定启动预加载专家数（默认自动封顶 4096）
--ssd-streaming-cold冷启动（不预加载 hotlist；仅用于测量最坏情况）
--simulate-used-memory SIZE模拟已占用内存

从用户视角看，SSD streaming 是以 Metal 为主的容量模式（上游 README 措辞为 "Metal-only capacity mode"）。但 streaming expert cache 的抽象层在 CUDA/ROCm 后端同样已实现（见上文 GPU 后端集成），只是成熟度和文档化程度不同。

相关概念 ​

  • moe — MoE 路由专家是 SSD streaming 的缓存对象
  • mmap — SSD streaming 使用 mmap 按需加载专家权重
  • quantization — 混合精度路由专家（per-layer boosted quants）触发 streaming 的特殊处理路径
  • distributed-inference — 另一种突破单机内存限制的方式

详见 Part 6 — GPU 加速。

出链