| """Selective-token forward for the block-causal Wan DiT. |
| |
| MotionCache (token granularity) and FlowCache (frame-group granularity) both need |
| to run the transformer stack on a *subset* of the current chunk's tokens while the |
| remaining tokens keep the activations cached from the previous computed step. |
| |
| For a KV-cached causal model that means, per self-attention layer: |
| * compute q/k/v only for the selected rows, |
| * scatter the fresh k/v into the slots the current chunk owns in the KV cache, |
| leaving the unselected slots holding the previous step's k/v, |
| * attend the selected queries against the *full* cache (history + current chunk). |
| |
| This mirrors ``forward_selective`` in |
| ``MotionCache4SkyReels-V2/skyreels_v2_infer/modules/transformer_motion.py``, |
| adapted from SkyReels' persistent ``k_cache/v_cache`` buffers to Self-Forcing's |
| rolling ``kv_cache[block]`` dict. |
| """ |
|
|
| import math |
|
|
| import torch |
|
|
| from wan.modules.attention import attention |
|
|
|
|
| def build_causal_freqs(grid_sizes, freqs, start_frame): |
| """Return the per-token RoPE multipliers [seq_len, 1, c] for one chunk. |
| |
| Same construction as ``causal_rope_apply`` in ``wan/modules/causal_model.py``, |
| factored out so it can be indexed by a token subset and reused across layers. |
| """ |
| c = freqs.size(1) |
| parts = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) |
| f, h, w = grid_sizes[0].tolist() |
| seq_len = f * h * w |
| freqs_i = torch.cat( |
| [ |
| parts[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), |
| parts[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), |
| parts[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), |
| ], |
| dim=-1, |
| ).reshape(seq_len, 1, -1) |
| return freqs_i |
|
|
|
|
| def rope_apply_rows(x, freqs_rows): |
| """Apply precomputed RoPE multipliers to x [B, S, n, d] given freqs [S, 1, d/2].""" |
| b, s, n, _ = x.shape |
| x_c = torch.view_as_complex(x.to(torch.float64).reshape(b, s, n, -1, 2)) |
| out = torch.view_as_real(x_c * freqs_rows.unsqueeze(0)).flatten(3) |
| return out.type_as(x) |
|
|
|
|
| def self_attn_selective(attn, x_sel, sel_idx, freqs_rows, kv_cache, |
| current_start, block_seq_len): |
| """Self-attention over a token subset, refreshing only those KV slots. |
| |
| Args: |
| attn: a ``CausalWanSelfAttention`` module. |
| x_sel: [B, n_sel, C] modulated hidden states for the selected tokens. |
| sel_idx: LongTensor [n_sel], token offsets inside the current chunk. |
| freqs_rows: [block_seq_len, 1, d/2] RoPE multipliers for the whole chunk. |
| kv_cache: this layer's ``{"k", "v", "global_end_index", "local_end_index"}``. |
| current_start: token offset of the chunk in global sequence coordinates. |
| block_seq_len: number of tokens the chunk owns (all frames of the chunk). |
| """ |
| b, n_sel = x_sel.shape[:2] |
| n, d = attn.num_heads, attn.head_dim |
|
|
| q = attn.norm_q(attn.q(x_sel)).view(b, n_sel, n, d) |
| k = attn.norm_k(attn.k(x_sel)).view(b, n_sel, n, d) |
| v = attn.v(x_sel).view(b, n_sel, n, d) |
|
|
| rows = freqs_rows[sel_idx] |
| roped_q = rope_apply_rows(q, rows) |
| roped_k = rope_apply_rows(k, rows).type_as(v) |
|
|
| |
| |
| current_end = current_start + block_seq_len |
| local_end_index = (kv_cache["local_end_index"].item() + current_end |
| - kv_cache["global_end_index"].item()) |
| local_start_index = local_end_index - block_seq_len |
|
|
| slots = local_start_index + sel_idx |
| kv_cache["k"].index_copy_(1, slots, roped_k) |
| kv_cache["v"].index_copy_(1, slots, v) |
|
|
| lo = max(0, local_end_index - attn.max_attention_size) |
| out = attention(roped_q, kv_cache["k"][:, lo:local_end_index], |
| kv_cache["v"][:, lo:local_end_index]) |
|
|
| kv_cache["global_end_index"].fill_(current_end) |
| kv_cache["local_end_index"].fill_(local_end_index) |
|
|
| return attn.o(out.flatten(2)) |
|
|
|
|
| def block_forward_selective(block, x_sel, sel_idx, frame_of_sel, e0, freqs_rows, |
| context, kv_cache, crossattn_cache, |
| current_start, block_seq_len): |
| """One ``CausalWanAttentionBlock`` evaluated on a token subset. |
| |
| ``frame_of_sel`` maps each selected token to its frame index inside the chunk, |
| so the per-frame AdaLN modulation can be gathered for the selected rows only. |
| """ |
| |
| e = (block.modulation.unsqueeze(1) + e0).chunk(6, dim=2) |
| e = [t.squeeze(2)[:, frame_of_sel] for t in e] |
|
|
| y = self_attn_selective( |
| block.self_attn, |
| block.norm1(x_sel) * (1 + e[1]) + e[0], |
| sel_idx, freqs_rows, kv_cache, current_start, block_seq_len, |
| ) |
| x_sel = x_sel + y * e[2] |
|
|
| x_sel = x_sel + block.cross_attn(block.norm3(x_sel), context, None, |
| crossattn_cache=crossattn_cache) |
| y = block.ffn(block.norm2(x_sel) * (1 + e[4]) + e[3]) |
| x_sel = x_sel + y * e[5] |
| return x_sel |
|
|
|
|
| def run_blocks_selective(model, x, sel_idx, e0, kwargs, kv_cache, crossattn_cache, |
| current_start, cache_start, grid_sizes): |
| """Run all transformer blocks on ``x[:, sel_idx]`` and return the updated rows. |
| |
| Returns a tensor [B, n_sel, C] holding the new hidden states for the selected |
| tokens (the caller merges them into the cached full-length tensor). |
| """ |
| f, h, w = grid_sizes[0].tolist() |
| frame_seqlen = h * w |
| block_seq_len = f * frame_seqlen |
| frame_of_sel = torch.div(sel_idx, frame_seqlen, rounding_mode="floor") |
|
|
| freqs_rows = build_causal_freqs(grid_sizes, kwargs["freqs"], |
| current_start // frame_seqlen) |
|
|
| x_sel = x[:, sel_idx] |
| for i, block in enumerate(model.blocks): |
| x_sel = block_forward_selective( |
| block, x_sel, sel_idx, frame_of_sel, e0, freqs_rows, |
| kwargs["context"], kv_cache[i], crossattn_cache[i], |
| current_start, block_seq_len, |
| ) |
| return x_sel |
|
|
