解构 DeepSeek 的 JIT：DeepJIT

发表于 2026-10-21 分类于 AI ， Architect

拆开 DeepSeek 开源的 JIT 运行时 DeepJIT，看一个 2770 行的头文件库如何用内容寻址缓存和整目录 rename 解决多机复用编译产物，又在进程内契约上留下了哪些松口。

GPU 上的 kernel 不是编译一次就完事的。同一个算子，shape 变了、dtype 变了、硬件换了、编译选项调了，都得重新生成一份二进制。如果每次启动都现场编译，服务冷启动的那几分钟就在等编译器。

DeepSeek 把这个环节单独抽出来做成了一个库：DeepJIT。定位很窄——只为在运行期编译、缓存、加载、启动 device kernel 提供基础设施，里面没有任何具体算子。

窄有窄的好处。2770 行主体代码，把一件事情的每个角落都想清楚了；窄也有窄的代价，它必须假设调用方在边界上是守规矩的。这篇就沿着这两条线去拆。

一、形态：一个塞进别人扩展里的头文件库

header-only，C++20，主体 2770 行，其中工具层约 38%、后端约 31%、运行时与缓存约 12%。
双后端：CUDA（NVCC 出 CUBIN，Driver API 加载）与 Ascend（bisheng 编译、ld.lld 链接，ACL 加载）。
嵌入形态：作为头文件被上层 kernel library 编进自己的扩展，进程内运行，通过 pybind11 暴露一个 get_jit()。
测试资产比库本身还大：C++ 4356 行、Python 1367 行。

一个 header-only 的库，意味着消费方不需要链接它、不需要处理 ABI、不需要在部署时多带一个 .so。代价是所有实现都在头文件里，编译时间、符号膨胀、模板实例化深度全部转嫁给使用者。

这也解释了它后面很多设计选择：为什么用模板而不是虚接口，为什么大量用宏做懒绑定，为什么对外只有一个极小的绑定面。

二、骨架：五层，一次静态多态

graph TB
  subgraph 绑定层
    A[python_api / GIL 守卫]
  end
  subgraph 运行时层
    B["Runtime 模板（参数化后端）"]
    C[Config / Env]
    D[Parser: 递归 include 哈希]
  end
  subgraph 缓存层
    E[MemCache: 句柄]
    F[DiskCache: 原子发布]
  end
  subgraph 后端层
    G[CUDA]
    H[Ascend]
  end
  subgraph 工具层
    I["exception / hash / env / command / fs / lazy"]
  end
  A --> B
  B --> C
  B --> D
  B --> E
  E --> F
  B --> G
  B --> H
  G --> I
  H --> I

依赖严格单向，工具层完全不感知上层。唯一一处瑕疵在文档里被点名：后端为了拿 Config 的定义，include 了 runtime/runtime.hpp 而不是 runtime/config.hpp。当前不成环，但只要哪天 runtime.hpp 反过来引入后端头，头文件环立刻出现。

真正决定这套结构气质的，是 Runtime<Backend> 这个模板。两个后端不是虚基类的两个实现，而是同一个模板的两个实例：

零虚调用开销，对每个 kernel launch 都是直接调用；
编译期裁剪，用 CUDA 的人根本不会编进 Ascend 的代码；
后端差异全部关在目录里，共享层没有一处 #ifdef。

代价是接口靠约定：后端需要提供哪些类型别名、哪些静态成员、哪些成员函数，全靠文档和编译器在实例化深处报错。加一个 concept Backend 能把报错提前到定义处，但那是后话。

三、主链路：一次 compile 要走几道门

flowchart LR
  A[compile] --> B[合并选项]
  B --> C[算 cache key]
  C --> D{内存缓存}
  D -- 命中 --> Z[shared_ptr Kernel]
  D -- 未命中 --> E{磁盘缓存}
  E -- 命中 --> H[加载]
  E -- 未命中 --> F[调用编译器]
  F --> G[原子发布]
  G --> H
  H --> I[写回内存缓存]
  I --> Z
  Z --> J[launch]

这条链上没有一步是多余的。值得逐个看的，是每一道门的判据。

第一道门：缓存键。 键由五段信息链式混进同一个 FNV1a 状态：应用提供的扩展签名、编译器版本、最终编译选项、post hook 脚本内容、源码加递归 include 的头文件内容。任何一段变化，就是一个全新的 128 位键。

这里有个很克制的判断：参与哈希的是”最终 flags 串”，不是选项字段。选项字段和命令行之间是多对一关系，比如打开 debug 会派生出好几项 flag；如果哈希字段，重构一次字段命名就让全世界的缓存失效；哈希最终命令，只有真正改变了编译行为才失效。反过来的代价是 flag 的组装顺序必须稳定，顺序一动，缓存全灭。

第二道门：内存缓存。 一个 unordered_map<string, shared_ptr<Kernel>>，命中的是已经加载好的设备句柄。用 shared_ptr 而不是裸句柄，是为了让跨 Runtime 实例的共享也能靠引用计数延后卸载。

第三道门：磁盘缓存。 命中就加载，未命中才真正调编译器。编译完成后不是直接写目标位置，而是走一套发布协议——这是整个库最值得读的部分。

四、磁盘缓存：用文件系统原语做分布式协调

需求很具体：多个用户、多个进程、多个节点共享一个缓存目录，复用彼此编译出来的 kernel。

书上的做法是加锁。但分布式文件系统的锁语义各家不一，一旦依赖它，可移植性就没了。DeepJIT 选了另一条路：不锁，允许重复计算，只保证”发布即完整”。

sequenceDiagram
  participant P1 as 进程 A
  participant FS as 共享文件系统
  participant P2 as 进程 B
  P1->>FS: 建 tmp/uuid 临时目录
  P1->>FS: 编译产物写入临时目录
  P1->>FS: 写 .committed 标记并 fsync
  P1->>FS: 自底向上 fsync 目录树
  P1->>FS: 整目录 rename 到 cache/tag.digest
  P2->>FS: 同样编译，rename 时发现已存在
  P2->>FS: 清理自己的临时目录，采用既有结果

四步里每一步都有理由：

.committed 标记先落盘并 fsync，保证”目录里有标记”就意味着内容齐全；
自底向上 fsync 整棵树，保证 rename 之后别人看到的不是半截数据；
整目录 rename 是 POSIX 上少有的、普遍原子的目录级发布原语；
rename 失败不当作错误，而是”别人先到了”，清理自己的临时目录，用对方的结果。

于是读者只有两种视角：看不到这个目录，或者看到一个完整目录。中间态永远不对外可见。这套协议不需要任何锁，也不需要任何元数据服务，只要文件系统满足”同一文件系统内 rename 原子 + fsync 生效”。

同一个机制还顺带支持了多根查找：DJ_JIT_CACHE_DIR 可以配成 $HOME/.dj:/shared/deep_jit，前者个人可写，后者团队只读，查找按顺序走，未命中只写第一个根。个人缓存和团队缓存就这么叠起来了。

还有一个细节值得记：命中时会顺手刷一下目录的 mtime，当作 LRU 的记账。不是精确 LRU，但成本几乎为零，够用。

五、双后端：同一套流程，两种世界观

CUDA 侧比较直：写 kernel.cu，拼 nvcc --cubin，校验 ptxas 的警告（寄存器溢出、local memory），可选跑一个 Python post hook、dump PTX/SASS，最后写 meta.json 记录命令与配置。

Ascend 侧多了一道工序：bisheng 先出可重定位目标，ld.lld 再链接成最终产物，中间文件删掉。加载时不能像 CUDA 那样直接从 CUBIN 里问出 kernel 名，得手写 ELF 解析，从 .ascend.meta.<name> 这样的节名里把名字提出来，还要求唯一。启动参数也不是简单打包，要按 ABI 对齐规则塞进连续缓冲区，最小 4 字节对齐。

这两套东西的差异之大，正好说明了模板参数化的价值：调用方写的是同一个 compile / launch，背后是两家完全不同的工具链、加载方式和参数约定。差异被关在两个目录里，共享层一行条件编译都没有。

代价也很诚实：两边的代码是逐文件同构的，同步靠人。重复度工具给出的 1.7% 很漂亮，但那是文本重复率——选项定义在四个地方手写同步（默认值、覆盖、flags 生成、JSON 序列化），这种语义重复工具测不出来，下一次加字段照样漏。

六、三条可以直接抄的设计

第一，内容寻址 + 原子发布，把并发问题变成幂等问题。 缓存条目的名字就是输入的摘要，同样的输入必然得到同一个名字，于是”谁先写”根本不重要，重复计算只是浪费一点 CPU，不会产生错误结果。这比加锁简单得多，也比分片、租约、协调服务好维护。

第二，允许失败者做无用功。 协议里没有”抢占”这个概念，两个进程同时编译同一个 kernel 是正常现象，后到的那个认输即可。工程上很多并发设计之所以复杂，就是因为不肯接受这点浪费。

第三，RAII 加不可变句柄。 kernel 句柄删除拷贝与移动，只能通过 shared_ptr 共享，析构里 unload 且保证不抛。异常路径上的资源泄漏被析构兜住了，这是教科书式的正确，也是嵌入式库该有的样子。

顺带说一句，错误模型选了异常而不是错误码：错误路径极少发生，异常让主路径零负担，还能带上文件行号和回溯。代价是库不能用在禁异常的环境里，这条在文档里写明了，算是把边界讲清楚。

七、三处松口

如果只看上面，会觉得这是个挑不出毛病的库。但它的松劲全在进程内，而且松得很有规律：跨进程的竞争被认真对待了，进程内的共享状态却没有主人。

其一，三处永不失效的缓存。 内存缓存无锁、无上限；include 解析的缓存一旦建立就不再失效；post hook 的内容按线程各缓存一份。后果是：在一个长驻进程里，如果有人改了头文件，缓存键可能还是旧的，于是拿到一份静默过期的二进制——不报错，只是算出来的结果不对。对一个推理系统来说，这是最难查的一类故障。

更麻烦的是键完备性的责任划分：库只跟踪 include_prefixes 命中的头文件，没命中的依赖要靠调用方自己往 extra_signature 里塞。这个设计本身合理，但它没有被写成一条硬契约，而是留在了文档的角落里。

其二，可信边界没有声明。 tag 做了白名单，防住了路径注入；但缓存根、include 路径这些输入并没有做根目录校验，../ 可以逃出去。编译命令用 popen 拼接，路径里带空格或特殊字符就会出问题。加上几处直接往 stdout 写日志，对一个嵌在宿主进程里的库来说都不够干净。

这些问题的共同点是：它们都假设调用方是善意的、环境是干净的。在内网、由熟悉系统的团队使用时，这个假设成立；一旦这个库被当成通用组件分发，假设就不成立了。

其三，没有观测。 一个以”复用编译产物”为核心卖点的库，没有缓存命中率、没有编译耗时统计、没有 tmp 残留回收。出问题只能靠翻 meta.json 和试环境变量。

对比之下，测试反而是它的亮点：C++ 测试 4356 行、Python 测试 1367 行，两倍于库本身，还包含负例和并发协议的自测。但这些测试锁在”有 GPU 或 NPU 的机器上才能跑”的保险柜里，没有 CI 就约等于没有门禁。

八、怎么用这份设计

如果你要解决的问题是”同一个进程内，按需编译 kernel 并复用”，DeepJIT 的缓存协议可以直接照抄：内容寻址、两阶段提交、整目录 rename、多根查找，这一套在本地文件系统上同样成立，而且比任何自研的锁方案都好维护。

如果要解决的是”多机共享编译产物”，那要先确认三件事：文件系统是否提供同 FS 的原子 rename 和 fsync、缓存目录是否可信、键是否真的覆盖了全部输入。前两条是环境假设，第三条是自己的责任——键不完备的缓存，比没有缓存更危险，因为它会把错误藏起来。

至于进程内那部分，我的看法是别照抄：内存缓存要加锁或明确声明单线程，include 缓存要带 mtime/size 失效，hook 缓存要么全局一份要么带失效，路径要收敛到库自己的根目录下。这些都是小改动，但决定了这个库能不能被放心地嵌进一个跑几个月不退出的进程。

一句话总结这个项目：把跨进程的协议写对了，把进程内的契约写松了。前者是功力，后者是注意力。

对一个 JIT 基础设施来说，缓存协议是它最难的部分，它做对了；缓存一致性是它最不能错的部分，它还欠着。这个顺序其实很常见——先把看得见的竞争解决好，再回头处理那些”默认调用方会小心”的地方。区别只在于，后者往往要等到一次线上事故才会被想起来。