Skip to content

fix(cache): shard Fn residual under context parallelism for CP + cache - #1113

Merged
DefTruth merged 1 commit into
mainfrom
dev
Sep 7, 2026
Merged

fix(cache): shard Fn residual under context parallelism for CP + cache#1113
DefTruth merged 1 commit into
mainfrom
dev

Conversation

@DefTruth

@DefTruth DefTruth commented Sep 7, 2026

Copy link
Copy Markdown
Member

fixed #1094

The cache Fn-blocks residual (_get_Fn_residual) subtracted a full-sequence original_hidden_states from the already CP-sharded hidden_states, crashing with a shape mismatch (e.g. 1280 vs 640) when --cache was combined with --parallel ulysses/ring (issue #1094). The CP-enabled guard only inspected diffusers' native _diffusers_hook registry, while cache-dit's own ulysses/ring runtime attaches split hooks to the vendored _cache_dit_hook registry (_ContextParallelSplitHook), so the guard was always False for the built-in CP backend and the shape-tolerance branch never triggered.

  • detect CP split hooks in BOTH diffusers (_diffusers_hook) and cache-dit (_cache_dit_hook) registries
  • shard the full-sequence original by replaying the block's own split hook (_prepare_cp_input) before subtracting, with an explicit shape guard and a proxy (hidden_states) fallback when no split hook matches
  • drop the now-unused HookRegistry import

Verified on Wan2.2-T2V: --cache --parallel ulysses|ring no longer crashes; real cache hits (49 steps) give ~2x speedup with structurally-identical output (no garbling); the non-CP path is byte-for-byte unchanged.

…ing + cache

The cache Fn-blocks residual (_get_Fn_residual) subtracted a full-sequence
original_hidden_states from the already CP-sharded hidden_states, crashing with
a shape mismatch (e.g. 1280 vs 640) when --cache was combined with
--parallel ulysses/ring (issue #1094). The CP-enabled guard only inspected
diffusers' native _diffusers_hook registry, while cache-dit's own ulysses/ring
runtime attaches split hooks to the vendored _cache_dit_hook registry
(_ContextParallelSplitHook), so the guard was always False for the built-in CP
backend and the shape-tolerance branch never triggered.

- detect CP split hooks in BOTH diffusers (_diffusers_hook) and cache-dit
  (_cache_dit_hook) registries
- shard the full-sequence original by replaying the block's own split hook
  (_prepare_cp_input) before subtracting, with an explicit shape guard and a
  proxy (hidden_states) fallback when no split hook matches
- drop the now-unused HookRegistry import

Verified on Wan2.2-T2V: --cache --parallel ulysses|ring no longer crashes;
real cache hits (49 steps) give ~2x speedup with structurally-identical output
(no garbling); the non-CP path is byte-for-byte unchanged.

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
@DefTruth
DefTruth merged commit 3db8d1e into main Sep 7, 2026
4 checks passed
@DefTruth
DefTruth deleted the dev branch September 7, 2026 04:46
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.

[Bug] --cache 与 --parallel ulysses 组合崩溃:_get_Fn_residual 中 hidden_states 维度不匹配(分片 vs 未分片)

1 participant