diff --git a/v2/platform/backends/torch_backend.py b/v2/platform/backends/torch_backend.py index 6aa2563a1..c404bd2f0 100644 --- a/v2/platform/backends/torch_backend.py +++ b/v2/platform/backends/torch_backend.py @@ -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: diff --git a/v2/recipes/wan_causal/loop.py b/v2/recipes/wan_causal/loop.py index 555d5e0bb..2fb61c2d9 100644 --- a/v2/recipes/wan_causal/loop.py +++ b/v2/recipes/wan_causal/loop.py @@ -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")