File size: 6,821 Bytes
b34c6c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Block-level re-implementation of ``CausalWanAttentionBlock.forward`` (LingBot
``wan/modules/model_fast.py``) in three flavours:
* ``block_forward`` -- the full block, optionally capturing the
self-attn / cross-attn / FFN outputs (TaylorSeer);
* ``block_forward_selective`` -- the block on a token subset, refreshing only the
selected KV slots (MotionCache);
* ``block_cam`` -- the per-token camera scale/shift a block applies
after self-attention (constant within a chunk).
Differences from ``cachelib/selective.py`` (Self-Forcing): the AdaLN modulation
``e`` is already per token ([B, L, 6, C]), the camera injection sits between the
self-attention residual and the cross-attention, and the self-attention takes
``max_attention_size`` / ``frame_seqlen`` from the call. Exactness against the
stock forward is checked by ``lb_test_exact.py``.
"""
import torch
import torch.nn.functional as torch_F
from wan.modules.attention import attention
def block_mod(block, e0):
"""The six per-token modulation tensors [B, L, C] of one block."""
with torch.amp.autocast('cuda', dtype=torch.float32):
e = (block.modulation.unsqueeze(0) + e0).chunk(6, dim=2)
return [t.squeeze(2) for t in e]
def block_cam(block, plucker):
"""(scale, shift) [B, L, C] of the block's camera injection, or (None, None)."""
if plucker is None:
return None, None
h = block.cam_injector_layer2(torch_F.silu(block.cam_injector_layer1(plucker)))
h = h + plucker
return block.cam_scale_layer(h), block.cam_shift_layer(h)
def block_forward(block, x, e0, kw, kv_cache, crossattn_cache, current_start, features=None):
"""Stock block forward, op for op. ``features`` (dict) receives the pre-gate
self-attn / cross-attn / FFN outputs under keys ``sa`` / ``ca`` / ``ffn``."""
e = block_mod(block, e0)
y = block.self_attn(
block.norm1(x).float() * (1 + e[1]) + e[0],
kw["seq_lens"], kw["grid_sizes"], kw["freqs"], kv_cache, current_start,
kw["max_attention_size"], frame_seqlen=kw["frame_seqlen"], seq_lens_int=None)
if features is not None:
features["sa"] = y
with torch.amp.autocast('cuda', dtype=torch.float32):
x = x + y * e[2]
plucker = (kw["dit_cond_dict"] or {}).get("c2ws_plucker_emb")
scale, shift = block_cam(block, plucker)
if scale is not None:
x = (1.0 + scale) * x + shift
y = block.cross_attn(block.norm3(x), kw["context"], kw["context_lens"],
crossattn_cache=crossattn_cache,
cross_attn_first_call=kw["cross_attn_first_call"])
if features is not None:
features["ca"] = y
x = x + y
y = block.ffn(block.norm2(x).float() * (1 + e[4]) + e[3])
if features is not None:
features["ffn"] = y
with torch.amp.autocast('cuda', dtype=torch.float32):
x = x + y * e[5]
return x
def build_causal_freqs(grid_sizes, freqs, start_frame):
"""Per-token RoPE multipliers [seq_len, 1, c] for one chunk (``causal_rope_apply``)."""
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
return 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)
def rope_apply_rows(x, freqs_rows):
b, s, n, _ = x.shape
x_c = torch.view_as_complex(x.to(torch.float64).reshape(b, s, n, -1, 2))
return torch.view_as_real(x_c * freqs_rows.unsqueeze(0)).flatten(3).type_as(x)
def self_attn_selective(attn, x_sel, sel_idx, freqs_rows, kv_cache, current_start,
block_seq_len, max_attention_size):
"""Self-attention over a token subset, refreshing only those KV slots. Only
valid on a repeat step of a chunk (its slots already exist: no eviction)."""
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).type_as(v)
roped_k = rope_apply_rows(k, rows).type_as(v)
current_end = current_start + block_seq_len
assert current_end == kv_cache["global_end_index"].item(), "selective step must repeat a chunk"
local_end_index = kv_cache["local_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 - max_attention_size)
out = attention(roped_q, kv_cache["k"][:, lo:local_end_index], kv_cache["v"][:, lo:local_end_index])
return attn.o(out.flatten(2))
def block_forward_selective(block, x_sel, sel_idx, e0, freqs_rows, kw, kv_cache,
crossattn_cache, current_start, block_seq_len):
e = [t[:, sel_idx] for t in block_mod(block, e0)]
y = self_attn_selective(
block.self_attn, block.norm1(x_sel).float() * (1 + e[1]) + e[0], sel_idx, freqs_rows,
kv_cache, current_start, block_seq_len, kw["max_attention_size"])
with torch.amp.autocast('cuda', dtype=torch.float32):
x_sel = x_sel + y * e[2]
plucker = (kw["dit_cond_dict"] or {}).get("c2ws_plucker_emb")
scale, shift = block_cam(block, None if plucker is None else plucker[:, sel_idx])
if scale is not None:
x_sel = (1.0 + scale) * x_sel + shift
y = block.cross_attn(block.norm3(x_sel), kw["context"], kw["context_lens"],
crossattn_cache=crossattn_cache, cross_attn_first_call=False)
x_sel = x_sel + y
y = block.ffn(block.norm2(x_sel).float() * (1 + e[4]) + e[3])
with torch.amp.autocast('cuda', dtype=torch.float32):
x_sel = x_sel + y * e[5]
return x_sel
def run_blocks_selective(model, x, sel_idx, e0, kw, kv_cache, crossattn_cache, current_start):
"""Run all blocks on ``x[:, sel_idx]``; returns the new rows [B, n_sel, C]."""
grid_sizes = kw["grid_sizes"]
f, h, w = grid_sizes[0].tolist()
frame_seqlen = h * w
block_seq_len = f * frame_seqlen
freqs_rows = build_causal_freqs(grid_sizes, kw["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, e0, freqs_rows, kw,
kv_cache[i], crossattn_cache[i], current_start, block_seq_len)
return x_sel
|