"""Blockcausal Wan forward shared by training and streaming inference. One implementation serves both regimes, which is the point: what the student is trained under is literally the code that runs at deployment. inference BufferKV -> preallocated StreamingKVCache, no_grad, batch 1 training TrainKV -> plain per-layer tensors, autograd-safe, batched The contract in both cases is the papers' "KV construction": the past is a *clean* (t=0) key/value prefix laid down by an earlier pass, and the current block is the only thing carrying noise. Block-causality is therefore structural -- the key set only ever holds past + current -- so no attention mask is needed over it, and adding one would wrongly serialise tokens within a block. Two ways to run the noisy blocks of a clip, both against the same clean prefix: block_forward one block, sequential. What inference does. parallel_blocks_forward every block of the clip as a batch row, one forward. Each row carries its own noise level and its own absolute RoPE offset, and reads the prefix slice its start index implies. Used for training. PER-FRAME TIMESTEP CONDITIONING (v0.2). Upstream Wan folds one shared `e0` into every token, so a sequence mixing a clean past with a noisy present is inexpressible in a single forward. `time_embed` here accepts a per-latent-frame vector, giving modulation [F, 6, dim] applied to x viewed as [B, F, S, dim]. With the pre-computed clean K/V above it is not *required*. The clean prefix is laid down by its own t=0 pass -- but it is what makes a block of mixed noise levels expressible at all, and a uniform scalar collapses to upstream's behaviour exactly, so the original weights load and run unchanged. """ import torch import torch.utils.checkpoint as ckpt from wan.modules.attention import flash_attention from wan.modules.model import sinusoidal_embedding_1d def time_embed(model, t, device, time_scale=1000.0): """Flow fraction(s) in [0,1] -> (e [N, dim], e0 [N, 6, dim]). N == 1 uniform timestep (upstream behaviour) N == batch rows one timestep per row N == latent frames per-frame timestep conditioning (batch must be 1) `time_scale` maps the internal [0,1] fraction onto the checkpoint's own convention: 1000 for stock Wan, 1.0 for a model fine-tuned on torch.rand(). """ if not torch.is_tensor(t): t = torch.as_tensor([float(t)], device=device) t = t.to(device=device, dtype=torch.float32).reshape(-1) # The time embedding runs outside autocast so its precision is the module's # own. Cast the input to match, rather than assuming fp32 weights: the # frozen teacher / critic base is held in bf16 to fit three networks on one # 40 GB card, and upstream's hard fp32 assumption would fail on it. wd = next(model.time_embedding.parameters()).dtype with torch.amp.autocast('cuda', enabled=False): e = model.time_embedding( sinusoidal_embedding_1d(model.freq_dim, t * time_scale).to(wd)) e0 = model.time_projection(e).unflatten(1, (6, model.dim)) return e, e0 class Modulation: """Per-layer AdaLN chunks for one set of timesteps. `(block.modulation + e0).chunk(6)` is identical for every token of a frame, so it is computed once per denoising step rather than per layer x frame. Each chunk is [N, 1, dim]. """ def __init__(self, blocks, e0): with torch.amp.autocast('cuda', enabled=False): self.chunks = [(blk.modulation + e0).chunk(6, dim=1) for blk in blocks] def __getitem__(self, i): return self.chunks[i] def _bcast(g, x, per_frame, n_frames): """Broadcast a modulation chunk [N,1,dim] against tokens x [B,L,dim]. per_frame=False: N is 1 or B, which already broadcasts. per_frame=True: N is the number of latent frames; tokens are frame-major, so view x as [B, F, S, dim] and give g a leading axis. """ if not per_frame: return g return g.unsqueeze(0) # [1, F, 1, dim] against [B, F, S, dim] def _as_frames(x, per_frame, n_frames): if not per_frame: return x b, l, d = x.shape return x.view(b, n_frames, l // n_frames, d) def _modulate(x, norm, shift, scale, per_frame, n_frames): y = _as_frames(norm(x).float(), per_frame, n_frames) y = y * (1 + _bcast(scale, x, per_frame, n_frames)) \ + _bcast(shift, x, per_frame, n_frames) return y.reshape(x.shape[0], x.shape[1], x.shape[2]) if per_frame else y def _gate(x, y, g, per_frame, n_frames): """x + y * g, in float32, with per-frame broadcasting.""" with torch.amp.autocast('cuda', dtype=torch.float32): if not per_frame: return x + y * g b, l, d = x.shape out = _as_frames(x, True, n_frames) + \ _as_frames(y, True, n_frames) * g.unsqueeze(0) return out.reshape(b, l, d) def _head_from(head, x, e, per_frame, n_frames): """Split out from `_head` so the FSDP shard unit can enter it with the head module it owns rather than reaching through the model (see wanstreamer.fsdp).""" with torch.amp.autocast('cuda', enabled=False): m = (head.modulation + e.unsqueeze(1)).chunk(2, dim=1) y = _modulate(x, head.norm, m[0], m[1], per_frame, n_frames) return head.head(y.to(head.head.weight.dtype)) def _head(model, x, e, per_frame, n_frames): return _head_from(model.head, x, e, per_frame, n_frames) def apply_rope(x, tbl): """x [B, L, n, d] real -> rotated, by tbl [L, 1, c] (shared) or [B, L, 1, c]. float32 complex rather than upstream's float64; verified equivalent within bf16 tolerance by tests/test_streaming_core.py. """ b, l, n, d = x.shape xc = torch.view_as_complex(x.float().reshape(b, l, n, d // 2, 2)) if tbl.dim() == 3: tbl = tbl.unsqueeze(0) return torch.view_as_real(xc * tbl).flatten(3) # ------------------------------------------------------------------- K/V stores class BufferKV: """Adapter over the preallocated StreamingKVCache used at inference.""" def __init__(self, cache): self.cache = cache def context(self, layer, k, v, **_): self.cache.write(layer, k, v) ck, cv = self.cache.context(layer, k.shape[1]) return ck, cv, None class BufferPrefixKV: """Read-only slice of a StreamingKVCache, in TrainKV's calling convention. The self-forcing trainer rolls out through the preallocated cache and then has to redo one recorded denoising step *with gradient*, against the prefix as it stood at that point. The buffer only ever grows during a rollout (no eviction is used in training), so `buffer[:upto]` is exactly that prefix, and reading it as a view keeps it detached for free. """ def __init__(self, cache): self.cache = cache def context(self, layer, k, v, upto=None, rows=1, k_lens=None): if not upto: return k, v, None pk = self.cache.k[layer, :upto].unsqueeze(0) pv = self.cache.v[layer, :upto].unsqueeze(0) if rows > 1: pk, pv = pk.expand(rows, -1, -1, -1), pv.expand(rows, -1, -1, -1) return torch.cat([k, pk], 1), torch.cat([v, pv], 1), k_lens class TrainKV: """Clean per-layer K/V prefix for a whole clip: plain tensors, autograd-safe. Built once per clip by `build_clean_kv` under no_grad; every noisy block then reads the slice its start index implies, so each block trains against exactly the prefix it would see at deployment. """ def __init__(self, num_layers): self.k = [None] * num_layers self.v = [None] * num_layers def append(self, layer, k, v): if self.k[layer] is None: self.k[layer], self.v[layer] = k, v else: self.k[layer] = torch.cat([self.k[layer], k], dim=1) self.v[layer] = torch.cat([self.v[layer], v], dim=1) @property def tokens(self): return 0 if self.k[0] is None else self.k[0].shape[1] def context(self, layer, k, v, upto=None, rows=1, k_lens=None): """Keys are laid out [current block ; clean prefix] -- current FIRST. That ordering is what makes the batched-block path expressible with a key-*length* mask: row i must see its own block plus prefix[:start_i*S], and only in this order are those two runs contiguous from index 0. Attention is permutation-invariant over keys, so the sequential path (which passes no mask at all) is unaffected. """ pk, pv = self.k[layer], self.v[layer] if pk is None or upto == 0: return k, v, None pk, pv = pk[:, :upto], pv[:, :upto] if rows > 1: pk = pk.expand(rows, -1, -1, -1) pv = pv.expand(rows, -1, -1, -1) return torch.cat([k, pk], 1), torch.cat([v, pv], 1), k_lens # ------------------------------------------------------------------- the layer def _layer(blk, x, ec, tbl, kv_ctx, ctx, ctx_lens, dtype, per_frame, n_frames): """One WanAttentionBlock in block-causal mode. kv_ctx(k,v) -> (K, V, k_lens).""" sa_in = _modulate(x, blk.norm1, ec[0], ec[1], per_frame, n_frames) b, s = sa_in.shape[0], sa_in.shape[1] n, d = blk.num_heads, blk.dim // blk.num_heads sa = blk.self_attn q = apply_rope(sa.norm_q(sa.q(sa_in)).view(b, s, n, d), tbl).to(dtype) k = apply_rope(sa.norm_k(sa.k(sa_in)).view(b, s, n, d), tbl).to(dtype) v = sa.v(sa_in).view(b, s, n, d).to(dtype) ck, cv, k_lens = kv_ctx(k, v) y = flash_attention(q=q, k=ck, v=cv, k_lens=k_lens, window_size=(-1, -1), causal=False) x = _gate(x, sa.o(y.flatten(2)), ec[2], per_frame, n_frames) x = x + blk.cross_attn(blk.norm3(x), ctx, ctx_lens) yf = blk.ffn(_modulate(x, blk.norm2, ec[3], ec[4], per_frame, n_frames)) return _gate(x, yf, ec[5], per_frame, n_frames) def _run(model, z, t, tbl, ctx, ctx_lens, kv_ctx_factory, dtype, time_scale, per_frame, grad_checkpoint, emb=None): """Shared body: patch-embed -> 30 block-causal layers -> head -> unpatchify. `emb` supplies a precomputed (e, e0). The CUDA-graph path needs it because `sinusoidal_embedding_1d` builds its frequency vector on the CPU and copies it to the device, which cannot be captured -- and the time embedding is two small linears, so hoisting it out of the graph costs nothing. """ B, _, F = z.shape[0], z.shape[1], z.shape[2] e, e0 = emb if emb is not None else time_embed(model, t, z.device, time_scale) nf = F if per_frame else 1 # `wanstreamer.fsdp.shard_model` attaches these: per-layer nn.Modules that # are the FSDP shard units. They must be ENTERED, because FSDP2 all-gathers # a module's parameters from a pre-forward hook on that module -- reaching # into `blk.self_attn.q` from outside, as `_layer` does, would run on # sharded parameters without ever erroring. Absent them nothing changes. layers = getattr(model, 'causal_layers', None) head_mod = getattr(model, 'causal_head', None) mod = Modulation(model.blocks, e0) if layers is None else None with torch.amp.autocast('cuda', dtype=dtype): x = model.patch_embedding(z.to(dtype)) gf, gh, gw = (int(s) for s in x.shape[2:]) x = x.flatten(2).transpose(1, 2) for li in range(len(model.blocks)): kv_ctx = kv_ctx_factory(li) if layers is None: fn = _layer args = (model.blocks[li], x, mod[li], tbl, kv_ctx, ctx, ctx_lens, dtype, per_frame, nf) else: # The sharded path passes e0 rather than a precomputed # modulation chunk: `blk.modulation` is a plain Parameter read # outside any forward by `Modulation`, so under FSDP it would # still be a shard at that point. CausalBlock folds that one add # into the layer, where the gather has already happened. It is # the same arithmetic and the same number of evaluations -- each # block's chunks are built exactly once per forward either way. fn = layers[li] args = (x, e0, tbl, kv_ctx, ctx, ctx_lens, dtype, per_frame, nf) if grad_checkpoint and torch.is_grad_enabled(): x = ckpt.checkpoint(fn, *args, use_reentrant=False) else: x = fn(*args) h = head_mod(x, e, per_frame, nf) if head_mod is not None \ else _head(model, x, e, per_frame, nf) return _unpatchify(model, h, gf, gh, gw) def _unpatchify(model, x, gf, gh, gw): """[B, L, out_dim*prod(patch)] -> [B, C, F, H, W]. Upstream's `WanModel.unpatchify` takes the grid as a *tensor* and calls `.tolist()` on it, which builds a CPU tensor and syncs -- neither is legal inside a CUDA graph capture. The grid is statically known here (it comes from the patch-embedding output shape), so do the same reshape with Python ints. Equivalence to upstream is asserted in scripts/verify_blockcausal.py. """ p0, p1, p2 = model.patch_size c = model.out_dim b = x.shape[0] u = x[:, :gf * gh * gw].view(b, gf, gh, gw, p0, p1, p2, c) u = u.permute(0, 7, 1, 4, 2, 5, 3, 6) # b c f p0 h p1 w p2 return u.reshape(b, c, gf * p0, gh * p1, gw * p2) # --------------------------------------------------------------- entry points def block_forward(model, z, t, t_start, rope, ctx, ctx_lens, kv=None, collect=None, dtype=torch.bfloat16, time_scale=1000.0, prefix_upto=None, per_frame=False, grad_checkpoint=False, tbl=None, emb=None): """Velocity for ONE block of latent frames against the clean prefix. z [B, C, F, H, W] noisy latents of the current block t float, or [F] flow fractions in [0,1] when per_frame t_start absolute latent-frame index of z[:, :, 0] (drives temporal RoPE) kv BufferKV (inference) | TrainKV (training) | None (self-attention only) collect list receiving this block's per-layer (k, v), or None tbl precomputed RoPE table; overrides t_start. The CUDA-graph path supplies it from a static buffer, since it is the one input that changes with the absolute frame index. """ if tbl is None: tbl = rope.span(t_start, z.shape[2]) def factory(li): def kv_ctx(k, v): if collect is not None: collect.append((k, v)) if kv is None: return k, v, None if prefix_upto is None: return kv.context(li, k, v) return kv.context(li, k, v, upto=prefix_upto) return kv_ctx return _run(model, z, t, tbl, ctx, ctx_lens, factory, dtype, time_scale, per_frame, grad_checkpoint, emb=emb) @torch.no_grad() def build_clean_kv(model, z, rope, ctx, ctx_lens, world_frames, block_frames, dtype=torch.bfloat16, time_scale=1000.0, rope_gap=0): """Clean (t=0) K/V for a whole clip, laid down exactly as deployment lays it down: the world primed as one bidirectional block, then each event block committed after being re-run at t=0. z: [B, C, F, H, W] clean latents. `rope_gap` shifts the *event* frames' temporal indices further from the world's -- see `event_rope_index`. """ kv = TrainKV(len(model.blocks)) F, S = z.shape[2], rope.seq spans, f = ([(0, world_frames)] if world_frames else []), world_frames while f < F: n = min(block_frames, F - f) spans.append((f, n)) f += n for t0, n in spans: collect = [] idx = t0 if t0 < world_frames or t0 == 0 else t0 + rope_gap block_forward(model, z[:, :, t0:t0 + n], 0.0, idx, rope, ctx, ctx_lens, kv=kv, collect=collect, dtype=dtype, time_scale=time_scale, prefix_upto=t0 * S) for li, (k, v) in enumerate(collect): kv.append(li, k, v) return kv def event_rope_index(frame, world_frames, rope_gap): """Absolute temporal index for RoPE, with the world pushed into the past. RoPE is relative: attention depends only on index differences. In a long stream the event window slides but the world block is *pinned*, so by unit 60 the current query sits ~180 latent frames from the world -- a relative distance training never showed the model if a clip is only 21 frames long, and rotary attention degrades badly off-distribution. Deployment produces that geometry for free (evicted events leave a real gap in the cache). Training has to simulate it, which costs nothing: keep the world at 0..world_frames and shift every event frame by a random gap. The content stays contiguous; only the positional distance changes, which is exactly the axis that needs covering. """ return frame if frame < world_frames else frame + rope_gap def parallel_blocks_forward(model, z_noisy, t_rows, starts, rope, ctx, ctx_lens, kv, block_frames, dtype=torch.bfloat16, time_scale=1000.0, grad_checkpoint=True): """Every event block of a clip in ONE forward, as batch rows. Row i holds block `starts[i]`, carries its own noise level `t_rows[i]` and its own absolute RoPE offset, and attends over [clean prefix[:starts[i]*S] ; its own noisy tokens]. Because the visible key set depends only on the row -- never on the query index -- this costs a [nb, 1, 1, Lk] key-padding mask instead of a dense [Lq, Lk] one (see sdpa_attention's key-padding fast path). z_noisy [1, C, F, H, W] full-clip noisy latents; returns [nb, C, b, H, W]. """ S, nb = rope.seq, len(starts) blocks = torch.cat([z_noisy[:, :, s:s + block_frames] for s in starts], 0) tbl = torch.stack([rope.span(s, block_frames).squeeze(1) for s in starts] ).unsqueeze(2) # [nb, b*S, 1, c] max_prefix = max(starts) * S # keys are [own block ; clean prefix], so row i keeps the first # block_frames*S + starts[i]*S of them (see TrainKV.context) k_lens = torch.tensor([block_frames * S + s * S for s in starts], device=z_noisy.device, dtype=torch.long) ctx_b = ctx.expand(nb, -1, -1) if ctx.shape[0] == 1 else ctx cl = ctx_lens.expand(nb) if ctx_lens.numel() == 1 else ctx_lens def factory(li): return lambda k, v: kv.context(li, k, v, upto=max_prefix, rows=nb, k_lens=k_lens) return _run(model, blocks, t_rows, tbl, ctx_b, cl, factory, dtype, time_scale, False, grad_checkpoint) def block_starts(num_frames, world_frames, block_frames): return [s for s in range(world_frames, num_frames, block_frames) if s + block_frames <= num_frames]