Skip to content

[Fix] Include world_size in cache directory name to prevent EP cross-contamination#42

Open
cennn wants to merge 2 commits into
SandAI-org:mainfrom
cennn:fix/world-size-cache-key
Open

[Fix] Include world_size in cache directory name to prevent EP cross-contamination#42
cennn wants to merge 2 commits into
SandAI-org:mainfrom
cennn:fix/world-size-cache-key

Conversation

@cennn

@cennn cennn commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • model_rank_dir_name() 的 cache 路径从 model_{idx}[_{tag}]_rank_{rank} 改为 model_{idx}[_{tag}]_rank_{rank}_ws{world_size}
  • EP=6 和 EP=8 切换时不再复用 stale 编译产物,避免 assert_size_stride 失败
  • 新增 3 个 cache invariant 回归测试

背景

切换 EP 配置(如 EP=6 → EP=8)时,MoE head padding 和 custom-op meta function 的 tensor stride 不同,但旧 cache key 不含 world_size,导致 rank 0 的编译产物被错误复用,运行时 assert_size_stride(expected_size != actual_size) 崩溃。

Test plan

  • 新增的 test_cache_invariant.py 3 个测试通过
  • 现有 CI 不受影响(仅新增 _ws{N} 后缀,对 world_size=1 单卡场景等价于 _ws1

jiahy0825
jiahy0825 previously approved these changes Jul 13, 2026

@jiahy0825 jiahy0825 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

…anager cache invariants

1. model_rank_dir_name() now produces `model_{idx}[_{tag}]_rank_{rank}_ws{world_size}`.
   Without world_size, switching between EP=6 and EP=8 silently reuses stale
   compiled artifacts whose tensor strides no longer match, causing
   assert_size_stride failures at runtime.

2. CompilerManager cache hardening (3 bugs fixed):
   - initialize_cache(): unconditionally reset self.cache and
     _remaining_restart_skips before loading index, preventing ghost
     handles from surviving across re-init calls
   - load(): check artifact directory exists on disk before delegating
     to compiler.load(); remove stale entry if missing
   - _maybe_store_cache_entry(): when handle is None (compilation failed),
     remove any pre-existing stale entry instead of silently keeping it

3. Add three cache-invariant regression tests reproducing the above bugs.
@cennn cennn force-pushed the fix/world-size-cache-key branch from 44369f9 to 780b160 Compare July 14, 2026 03:52
- Add unconditional _remaining_restart_skips reset in initialize_cache()
  to prevent stale skip counts from surviving across re-initialization
- Simplify test CompileConfig creation: use _cli_parse_args=False instead
  of sys.argv monkey-patching
- Add assertion verifying _remaining_restart_skips is cleared on re-init
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants