Make the JIT cache key independent of the install include path - #398
Make the JIT cache key independent of the install include path#398matteso1 wants to merge 2 commits into
Conversation
The key embedded the absolute -I<library include> token via flags, so byte-identical kernels re-JIT and duplicate cache entries whenever the install prefix differs. That directory's header content is already keyed through the include hash embedded in the code string; the compiler version signature, every behavioral flag, and every other include path (CUDA headers are not content-hashed) stay in the key. Compile commands are unchanged; only the key computation differs. Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
| // not kernel identity, as its headers' content is already hashed into `code`. | ||
| // Every other flag stays keyed, including other `-I` paths (e.g. the CUDA home | ||
| // include, whose headers are not content-hashed) and the compiler `signature`. | ||
| const auto key_flags = strip_flag_token(flags, "-I" + library_include_path.string()); |
There was a problem hiding this comment.
🔴 critical: 将捆绑的 CUTE/CUTLASS 头文件纳入缓存键: 当共享 DG_JIT_CACHE_DIR 的两个安装包含不同版本的 CUTE/CUTLASS 时，此处会产生相同缓存键并复用旧 CUBIN。setup.py:154-165 会把这些头文件复制到 library_include_path，DeepGEMM 头文件也直接包含它们，但 IncludeParser::get_includes 仅递归哈希 deep_gemm/*（include_parser.hpp:28）。因此在移除该 -I 标志前，需要将这些捆绑头文件的内容纳入键，否则更新它们后可能执行由旧安装编译的内核。
🤖 v6
There was a problem hiding this comment.
Verified this against current main. setup.py copies CUTE and CUTLASS into the installed include tree, while IncludeParser only followed deep_gemm/*, so the old diff could reuse a CUBIN across different bundled headers.
Fixed in 06a6ae2 by hashing every installed include file by relative path and content once per process, then including that tree hash in the generated-code hash. A local CPU regression now verifies that relocating identical trees is a hit and changing a bundled CUTE header is a miss. I also corrected the token-removal comment.
| fsync_path(path); | ||
| } | ||
|
|
||
| // Remove a single exact token from a whitespace-separated flag string |
There was a problem hiding this comment.
🔵 suggestion: 注释 "Remove a single exact token" 实际移除所有出现的该 token（当前 flags 中仅出现一次，无实际影响）。建议措辞改为 "Remove every exact token occurrence from a whitespace-separated flag string" 以更精确。
🤖 v4f
| fsync_path(path); | ||
| } | ||
|
|
||
| // Remove a single exact token from a whitespace-separated flag string |
There was a problem hiding this comment.
🔵 suggestion: 可顺带在注释中说明：若 include 路径本身含空白，strip 按空白分词将无法匹配该 token（key 退化为旧行为），且此类路径下 NVRTC 选项切分与 NVCC shell 命令本就会编译失败，属既有限制而非本 PR 回归。
🤖 v4f
| fsync_path(path); | ||
| } | ||
|
|
||
| // Remove a single exact token from a whitespace-separated flag string |
There was a problem hiding this comment.
🔵 suggestion: 注释写的是 “Remove a single exact token”，但循环实现实际会移除该 token 的所有出现位置。实践中该 token 只出现一次，且移除全部出现在语义上反而更安全，故不阻塞合并；建议把注释改为 “Remove all exact occurrences of a token” 以与实现一致。
🤖 v5
🤖 ds-review-bot Code Reviewv6该变更遗漏了同一包含目录下未被递归哈希的捆绑头文件，可能导致不同安装错误共享已编译内核。 v4f该 PR 正确且最小化地解决了 JIT 缓存键嵌入绝对 -I{library_include_path} 导致的跨安装前缀（改名 venv、容器构建路径 != 运行路径）无法共享 DG_JIT_CACHE_DIR 的问题。改动仅限 csrc/jit/compiler.hpp 一个文件：新增 strip_flag_token 辅助函数，build() 中仅对缓存键计算使用去除该 token 后的 key_flags；实际编译命令、DG_JIT_DEBUG 输出均未改动。未发现阻塞性问题。 v5审查通过（Approve）。该提交（865b0ce，重提 #388）将 JIT 缓存键与安装 include 路径解耦：在 Compiler::build() 计算缓存键时，仅剥离与 Files reviewed: 1 |
Signed-off-by: Nils Matteson <nilsmatteson@icloud.com>
|
|
||
| // Key the installed header tree by relative path and content, not its install prefix. | ||
| std::vector<std::filesystem::path> files; | ||
| for (const auto& entry: std::filesystem::recursive_directory_iterator(library_include_path)) { |
There was a problem hiding this comment.
Could this reintroduce the std::filesystem::recursive_directory_iterator ABI breakage addressed in 436a563? That change replaced this iterator with recursive directory_iterator traversal after it broke shipped wheels. Is the same compatibility concern applicable here?
|
|
||
| // Key the installed header tree by relative path and content, not its install prefix. | ||
| std::vector<std::filesystem::path> files; | ||
| for (const auto& entry: std::filesystem::recursive_directory_iterator(library_include_path)) { |
There was a problem hiding this comment.
develop.sh installs cute and cutlass as directory symlinks, while recursive_directory_iterator does not follow directory symlinks by default. Could that omit bundled headers from this hash and let source builds with different CUTE/CUTLASS revisions share a cache key and load a stale CUBIN?
3 participants
This resubmits #388 so the repository's current automated review runs.
The bug
The kernel cache key embeds the absolute
-I{library_include_path}token throughflags. Byte-identical kernels therefore re-JIT and create duplicate cache entries when the install prefix changes, such as a renamed virtual environment or a container build path that differs from the runtime path. That prevents a sharedDG_JIT_CACHE_DIRfrom being portable across those layouts.This follows simon-mo's request on vllm-project/vllm#48190 to propose the relocation fix upstream rather than carry it as a vendored patch. Related but distinct work includes #301/#302 for JIT cache concurrency and #333 for wheel portability.
The fix
Strip only the exact
-I{library_include_path}token when computing the cache key. The actual compile command andDG_JIT_DEBUGoutput are unchanged.To keep reuse safe across different installed bundles, the full installed include tree is hashed once per process by relative path and file content. That tree hash covers the packaged DeepGEMM, CUTE, and CUTLASS headers without depending on the install prefix. The compiler-version
signature, every behavioral flag, and every other include path also remain keyed. Existing entries re-JIT once under the new key.Measured
The broader vendored variant in vllm-project/vllm#48190 also drops the compiler-version field. On one H100 NVL with Qwen3-30B-A3B-FP8, TP1, and DeepGEMM
a6b593d, relocating a planted cache caused the stock key to recompile 41/41 kernels and start in 170.6 seconds. The tested portable key reused 41/41 and started in 76.9 seconds, a 93.7-second difference, n=1 per arm.This PR is intentionally narrower because it retains compiler-version identity. Its exact diff has not been independently timed on GPU hardware.
Test plan
git diff --checkmainat559d79fb6994a58b8a15b4b93bf13ccc16edf247