nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[bugfix] v2: SF-Wan cross-chunk KV cache — condition each chunk on prior clean chunks

Browse Source 
The causal adapter never passed a kv_cache, so CausalWanTransformer3DModel.forward routed to
_forward_train (no cross-chunk KV) on every chunk instead of _forward_inference (the CausVid
Alg-2 KV-cache path). Each chunk denoised blind to the previous ones; the loop's cross-chunk
"context" was a toy mean(prior_latents) the adapter ignored. Result: hard discontinuities at
every chunk boundary (frame-to-frame absdiff spikes 43-55 every ~12 frames).

Fix (cuda path only; toy/CPU path and the 237-test suite untouched):
- WanDiT.alloc_causal_caches(): allocate the persistent per-block KV + cross-attn caches sized
  from the model config (mirrors CausalDenoisingStage._initialize_kv_cache).
- WanDiT.__call__: thread kv_cache/crossattn_cache/current_start/cache_start/start_frame/
  frame_seqlen so the model runs _forward_inference.
- wan_causal/loop.py: own the caches in LoopState (per-request -> interleave-safe); pass
  current_start = chunk_idx*chunk_size*frame_seqlen per chunk; do the clean-KV write
  (timestep ~0) after each chunk so the next attends to it.

Verified on H100: frame-to-frame absdiff mean 12.8->4.4, max 55.4->8.3; chunk-boundary spikes
eliminated; coherent across all 7 chunks. CPU causal toy tests unchanged (24 passed).
This commit is contained in:
SolitaryThinker
2026-07-04 17:15:34 +00:00
committed by Peiyuan Zhang
parent 4ff33a28b7
commit f90e8f3e76
2 changed files with 69 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
v2/platform/backends/torch_backend.py
+40 -5
View File
@@ -207,8 +207,39 @@ class WanDiT(TorchComponent):
self.module.to(self.device)
self._on_gpu = True
def alloc_causal_caches(self, frame_seqlen: int, *, batch: int = 1, max_text_len: int = 512) -> Any:
"""Allocate the persistent per-block KV + cross-attn caches the causal student needs for
cross-chunk conditioning (mirrors CausalDenoisingStage._initialize_kv_cache). The chunk_rollout
loop owns these in its LoopState (per request, never on this shared adapter) and threads them
through __call__ so the model routes to _forward_inference."""
self._ensure_resident()
m = self.module
nblocks = len(m.blocks)
nheads, hdim = int(m.num_attention_heads), int(m.attention_head_dim)
local = int(getattr(m, "local_attn_size", -1) or -1)
if local != -1:
kv_size = local * frame_seqlen
else: # global window: model keeps the last GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES (21) frames
arch = getattr(getattr(m, "config", None), "arch_config", None)
sliding = int(getattr(arch, "sliding_window_num_frames", 0) or 0)
kv_size = max(sliding, 21) * frame_seqlen
dev, dt = self.device, self.dtype
kv = [{
"k": torch.zeros(batch, kv_size, nheads, hdim, device=dev, dtype=dt),
"v": torch.zeros(batch, kv_size, nheads, hdim, device=dev, dtype=dt),
"global_end_index": torch.tensor([0], dtype=torch.long, device=dev),
"local_end_index": torch.tensor([0], dtype=torch.long, device=dev),
} for _ in range(nblocks)]
ca = [{
"k": torch.zeros(batch, max_text_len, nheads, hdim, device=dev, dtype=dt),
"v": torch.zeros(batch, max_text_len, nheads, hdim, device=dev, dtype=dt),
"is_init": False,
} for _ in range(nblocks)]
return {"kv": kv, "crossattn": ca}
@torch.no_grad()
def __call__(self, latent, text_embed, sigma, context=None, *, cond=None):
def __call__(self, latent, text_embed, sigma, context=None, *, cond=None, kv_cache=None,
crossattn_cache=None, current_start=0, start_frame=0, frame_seqlen=1560):
self._ensure_resident()
hs = self._t(latent)
if cond is not None: # i2v: concat [noise (16ch) ; mask+cond_latent (20ch)] -> 36ch DiT input
@@ -219,11 +250,15 @@ class WanDiT(TorchComponent):
timestep = (torch.full(
(1, hs.shape[2]), ts, device=self.device) if self.causal else torch.tensor([ts], device=self.device))
img = None if self.causal else self._t(context) # i2v image embedding for standard Wan
fwd = dict(hidden_states=hs, encoder_hidden_states=ehs, timestep=timestep, encoder_hidden_states_image=img)
if self.causal and kv_cache is not None:
# Pass the persistent caches + frame offset so the model runs _forward_inference (CausVid
# Algorithm 2) and conditions this chunk on prior clean chunks. Without it the model falls
# through to _forward_train (no cross-chunk KV) -> chunks denoise blind -> discontinuity.
fwd.update(kv_cache=kv_cache, crossattn_cache=crossattn_cache, current_start=int(current_start),
cache_start=int(current_start), start_frame=int(start_frame), frame_seqlen=int(frame_seqlen))
with self._ctx():
velocity = self.module(hidden_states=hs,
encoder_hidden_states=ehs,
timestep=timestep,
encoder_hidden_states_image=img)
velocity = self.module(**fwd)
return self._n(velocity) # rectified-flow velocity (BRINGUP risk C)
def clone(self) -> WanDiT:
v2/recipes/wan_causal/loop.py
+29 -1
View File
@@ -75,6 +75,19 @@ class ChunkRolloutLoop:
caches=getattr(model, "caches", None))
st.latents["chunk"] = (rng.standard_normal(self._chunk_shape(req, model)) * float(sig[0])).astype("float32")
st.plugin_state["cfg"] = {}
# Real causal cross-chunk conditioning (cuda backend only): allocate the model's persistent
# KV + cross-attn caches and own them HERE in LoopState (per request -> interleave-safe), then
# thread them through every forward so the self-forcing student conditions each chunk on prior
# clean chunks. Toy/CPU keeps the slab mean-context path (``context`` above). See loop.next.
st.scratch["real_causal"] = False
dit = model.component("transformer")
if getattr(getattr(model, "platform", None), "device", "cpu") == "cuda" and hasattr(dit, "alloc_causal_caches"):
_, _, lh, lw = self._chunk_shape(req, model)
ps = getattr(getattr(dit, "module", None), "patch_size", (1, 2, 2))
frame_seqlen = (lh // int(ps[-2])) * (lw // int(ps[-1]))
caches = dit.alloc_causal_caches(frame_seqlen)
st.scratch.update(causal_kv=caches["kv"], causal_ca=caches["crossattn"],
frame_seqlen=frame_seqlen, real_causal=True)
return st
def next(self, st: LoopState):
@@ -90,6 +103,13 @@ class ChunkRolloutLoop:
scale = st.scratch["guidance_scale"]
context = np.mean(st.scratch["slabs"], axis=0) if st.scratch["slabs"] else None
cfg, precision = self.cfg, self.precision
# real-causal KV-cache threading (cuda backend); else the toy ``context`` path
real_causal = st.scratch.get("real_causal", False)
kv, ca = st.scratch.get("causal_kv"), st.scratch.get("causal_ca")
fsl = st.scratch.get("frame_seqlen", 1560)
ci = st.scratch["chunk_idx"]
cur_start, start_frame = ci * self.chunk_size * fsl, ci * self.chunk_size
is_last = (i == self.steps_per_chunk - 1)
def run(model, override=None):
fm = model.platform.kernels.get(FLOW_MATCH_STEP) # solver dispatched per (device, arch)
@@ -97,9 +117,17 @@ class ChunkRolloutLoop:
velocity = np.asarray(override["noise_pred"], dtype="float32")
else:
dit = model.component("transformer")
preds = {b: dit(x, pe if b == "cond" else ne, sigma_t, context=context) for b in branches}
if real_causal: # condition on prior clean chunks via the model's kv_cache (CausVid Alg 2)
kw = dict(kv_cache=kv, crossattn_cache=ca, current_start=cur_start,
start_frame=start_frame, frame_seqlen=fsl)
preds = {b: dit(x, pe if b == "cond" else ne, sigma_t, **kw) for b in branches}
else:
preds = {b: dit(x, pe if b == "cond" else ne, sigma_t, context=context) for b in branches}
velocity = cfg.combine(preds, scale, sctx, cfg_state)
x_next = fm(precision.cast(x), precision.cast(velocity), sigma_t, sigma_next)
if real_causal and is_last: # write the clean chunk's K/V (timestep ~0) so the next chunk attends to it
model.component("transformer")(x_next, pe, 0.0, kv_cache=kv, crossattn_cache=ca,
current_start=cur_start, start_frame=start_frame, frame_seqlen=fsl)
return StepResult(output={
"noise_pred": np.asarray(velocity, dtype="float32"),
"latents": x_next.astype("float32")