import math import types from types import SimpleNamespace import torch import torch.distributed as dist from wan.modules.attention import attention from wan.modules.model import sinusoidal_embedding_1d from wan.modules.causal_model_infinity import CausalWanModel as CausalWanModelInfinityBase from wan.modules.causal_model import CausalWanSelfAttention as _RelropeSelfAttention, causal_rope_apply _FIRST_LONG_LOGGED = {'flag': False} _FIRST_BULK_LOGGED = {'flag': False} def _self_attn_inf_mem_forward(self, x, seq_lens, grid_sizes, freqs, block_mask, kv_cache=None, current_start=0, cache_start=None, sink_recache_after_switch=False, memory_kv=None, capture_sink_qkv=False): b, s, n, d = (*x.shape[:2], self.num_heads, self.head_dim) if cache_start is None: cache_start = current_start def qkv_fn(_x): q = self.norm_q(self.q(_x)).view(b, s, n, d) k = self.norm_k(self.k(_x)).view(b, s, n, d) v = self.v(_x).view(b, s, n, d) return (q, k, v) q, k, v = qkv_fn(x) if kv_cache is None: raise NotImplementedError('CausalWanModelInfinityMemory only supports inference-time kv_cache forward; training path is intentionally removed.') frame_seqlen = math.prod(grid_sizes[0][1:]).item() num_new_frames = grid_sizes[0][0].item() current_end = current_start + q.shape[1] sink_tokens = self.sink_size * frame_seqlen sink_size_frames = self.sink_size kv_cache_size = kv_cache['k'].shape[1] num_new_tokens = q.shape[1] is_recompute = current_end <= kv_cache['global_end_index'].item() and current_start > 0 if self.local_attn_size != -1 and current_end > kv_cache['global_end_index'].item() and (num_new_tokens + kv_cache['local_end_index'].item() > kv_cache_size): num_evicted_tokens = num_new_tokens + kv_cache['local_end_index'].item() - kv_cache_size num_rolled_tokens = kv_cache['local_end_index'].item() - num_evicted_tokens - sink_tokens local_end_index = kv_cache['local_end_index'].item() + current_end - kv_cache['global_end_index'].item() - num_evicted_tokens local_start_index = local_end_index - num_new_tokens temp_k = kv_cache['k'].clone() temp_v = kv_cache['v'].clone() temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index roped_offset = max(0, write_start_index - local_start_index) write_len = max(0, local_end_index - write_start_index) if write_len > 0: temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len] temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len] cache_update_info = {'action': 'roll_and_insert', 'sink_tokens': sink_tokens, 'num_rolled_tokens': num_rolled_tokens, 'num_evicted_tokens': num_evicted_tokens, 'local_start_index': local_start_index, 'local_end_index': local_end_index, 'write_start_index': write_start_index, 'write_end_index': local_end_index, 'new_k': k[:, roped_offset:roped_offset + write_len], 'new_v': v[:, roped_offset:roped_offset + write_len], 'current_end': current_end, 'is_recompute': is_recompute} else: local_end_index = kv_cache['local_end_index'].item() + current_end - kv_cache['global_end_index'].item() local_start_index = local_end_index - num_new_tokens temp_k = kv_cache['k'].clone() temp_v = kv_cache['v'].clone() write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index if sink_recache_after_switch: write_start_index = local_start_index roped_offset = max(0, write_start_index - local_start_index) write_len = max(0, local_end_index - write_start_index) if write_len > 0: temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len] temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len] cache_update_info = {'action': 'direct_insert', 'local_start_index': local_start_index, 'local_end_index': local_end_index, 'write_start_index': write_start_index, 'write_end_index': local_end_index, 'new_k': k[:, roped_offset:roped_offset + write_len], 'new_v': v[:, roped_offset:roped_offset + write_len], 'current_end': current_end, 'is_recompute': is_recompute} num_cache_frames = local_end_index // frame_seqlen R_active = num_cache_frames - sink_size_frames N_Q_rr = memory_kv[0].shape[1] // frame_seqlen if memory_kv is not None else 0 current_start_frame = current_start // frame_seqlen num_frame_per_block_attr = getattr(self, 'num_frame_per_block_attr', 3) relative_rope_pmax = getattr(self, 'relative_rope_pmax', 21) rr_pos = _RelropeSelfAttention._compute_relative_positions(current_start_frame=current_start_frame, B=num_new_frames, R=R_active, N_Q=N_Q_rr, N_S=sink_size_frames, pmax=relative_rope_pmax, num_frame_per_block=num_frame_per_block_attr) _diag = f"[InfMem-diag] layer={getattr(self, '_layer_id', -1)} cur_start_f={current_start_frame} B={num_new_frames} R={R_active} N_Q={N_Q_rr} N_S={sink_size_frames} pmax={relative_rope_pmax} is_bulk={rr_pos['is_bulk_forward']} use_mem={rr_pos['use_memory']} sink_s={rr_pos['sink_start']} mem_s={rr_pos['mem_start']} mem_e={rr_pos['mem_end']} local_s={rr_pos['local_start']} local_e={rr_pos['local_end']} q_s={rr_pos['q_start']} q_l={rr_pos['q_last']}" assert rr_pos['q_last'] <= relative_rope_pmax - 1, f'InfMem overflow: q_last > pmax-1 | {_diag}' assert rr_pos['q_start'] >= 0, f'InfMem underflow: q_start < 0 | {_diag}' if rr_pos['use_memory']: assert rr_pos['mem_start'] >= sink_size_frames, f'InfMem mem overlaps sink | {_diag}' assert rr_pos['mem_end'] + 1 == rr_pos['local_start'], f'InfMem mem-local not contiguous | {_diag}' if R_active > 0 and (not rr_pos['is_bulk_forward']): assert rr_pos['q_last'] == rr_pos['local_end'], f'InfMem Q-local tail mismatch | {_diag}' if sink_tokens > 0: k_sink_raw = temp_k[:, :sink_tokens] v_sink = temp_v[:, :sink_tokens] sink_grid = grid_sizes.clone() sink_grid[:, 0] = sink_size_frames k_sink = causal_rope_apply(k_sink_raw, sink_grid, freqs, start_frame=0).type_as(v) else: k_sink = None v_sink = None if R_active > 0: k_local_raw = temp_k[:, sink_tokens:local_end_index] v_local = temp_v[:, sink_tokens:local_end_index] local_grid = grid_sizes.clone() local_grid[:, 0] = R_active k_local = causal_rope_apply(k_local_raw, local_grid, freqs, start_frame=rr_pos['local_start']).type_as(v) else: k_local = None v_local = None q_grid = grid_sizes.clone() q_grid[:, 0] = num_new_frames roped_query = causal_rope_apply(q, q_grid, freqs, start_frame=rr_pos['q_start']).type_as(v) mem_k_roped = None mem_v = None if memory_kv is not None and rr_pos['use_memory'] and (not rr_pos['is_bulk_forward']): mem_k_raw, mem_v_raw = memory_kv N_Q = mem_k_raw.shape[1] // frame_seqlen mem_grid = grid_sizes.clone() mem_grid[:, 0] = N_Q mem_k_roped = causal_rope_apply(mem_k_raw, mem_grid, freqs, start_frame=rr_pos['mem_start']).type_as(v) mem_v = mem_v_raw parts_k, parts_v = ([], []) if k_sink is not None: parts_k.append(k_sink) parts_v.append(v_sink) if mem_k_roped is not None: parts_k.append(mem_k_roped) parts_v.append(mem_v) if k_local is not None: parts_k.append(k_local) parts_v.append(v_local) k_cat = torch.cat(parts_k, dim=1) v_cat = torch.cat(parts_v, dim=1) _layer_id = getattr(self, '_layer_id', -1) if _layer_id == 0: if not rr_pos['is_bulk_forward'] and rr_pos['q_last'] == relative_rope_pmax - 1 and (not _FIRST_LONG_LOGGED['flag']): mem_str = f"mem[{rr_pos['mem_start']}..{rr_pos['mem_end']}]" if rr_pos['use_memory'] else 'mem=inactive' if not dist.is_initialized() or dist.get_rank() == 0: print(f"[InfMem] FIRST LONG PHASE @ t={current_start_frame}: sink[0..{sink_size_frames - 1}] | {mem_str} | local[{rr_pos['local_start']}..{rr_pos['local_end']}] | Q[{rr_pos['q_start']}..{rr_pos['q_last']}]", flush=True) _FIRST_LONG_LOGGED['flag'] = True if rr_pos['is_bulk_forward'] and (not _FIRST_BULK_LOGGED['flag']): if not dist.is_initialized() or dist.get_rank() == 0: print(f"[InfMem] FIRST BULK FORWARD @ t={current_start_frame}, B={num_new_frames}, R={R_active}: sink[0..{sink_size_frames - 1}] | gap | local[{rr_pos['local_start']}..{rr_pos['local_end']}] | Q[{rr_pos['q_start']}..{rr_pos['q_last']}] (memory skipped)", flush=True) _FIRST_BULK_LOGGED['flag'] = True x = attention(roped_query, k_cat, v_cat) x = x.flatten(2) x = self.o(x) return (x, (current_end, local_end_index, cache_update_info)) def _block_inf_mem_forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens, block_mask, kv_cache=None, crossattn_cache=None, current_start=0, cache_start=None, sink_recache_after_switch=False, memory_kv=None, capture_sink_qkv=False): num_frames, frame_seqlen = (e.shape[1], x.shape[1] // e.shape[1]) e_mod = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) self_attn_result = self.self_attn((self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e_mod[1]) + e_mod[0]).flatten(1, 2), seq_lens, grid_sizes, freqs, block_mask, kv_cache, current_start, cache_start, sink_recache_after_switch, memory_kv=memory_kv, capture_sink_qkv=capture_sink_qkv) if kv_cache is not None: y, cache_update_info = self_attn_result else: y = self_attn_result cache_update_info = None x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e_mod[2]).flatten(1, 2) x = x + self.cross_attn(self.norm3(x), context, context_lens, crossattn_cache=crossattn_cache) y = self.ffn((self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e_mod[4]) + e_mod[3]).flatten(1, 2)) x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e_mod[5]).flatten(1, 2) if cache_update_info is not None: return (x, cache_update_info) else: return x class CausalWanModelInfinityMemory(CausalWanModelInfinityBase): def enable_infmem(self, relative_rope: bool=True, relative_rope_pmax: int=21, num_frame_per_block_attr: int=3): self.relative_rope = relative_rope self.relative_rope_pmax = relative_rope_pmax self.num_frame_per_block_attr = num_frame_per_block_attr for i, block in enumerate(self.blocks): sa = block.self_attn sa.relative_rope = relative_rope sa.relative_rope_pmax = relative_rope_pmax sa.num_frame_per_block_attr = num_frame_per_block_attr sa._layer_id = i sa.forward = types.MethodType(_self_attn_inf_mem_forward, sa) block.forward = types.MethodType(_block_inf_mem_forward, block) object.__setattr__(self, 'query_memory_encoder', None) object.__setattr__(self, 'sink_memory', None) self._ei_prev_window_start = None _FIRST_LONG_LOGGED['flag'] = False _FIRST_BULK_LOGGED['flag'] = False def setup_memory_encoder(self, memory_kwargs): from model.query_memory import QueryMemoryEncoder if isinstance(memory_kwargs, dict): memory_kwargs = SimpleNamespace(**memory_kwargs) self.query_memory_encoder = QueryMemoryEncoder(memory_kwargs) @property def num_frame_per_block(self): return getattr(self, '_num_frame_per_block', 1) @num_frame_per_block.setter def num_frame_per_block(self, value): self._num_frame_per_block = int(value) for blk in self.blocks: blk.self_attn.num_frame_per_block_attr = int(value) def _forward_inference(self, x, t, context, seq_len, clip_fea=None, y=None, kv_cache: dict=None, crossattn_cache: dict=None, current_start: int=0, cache_start: int=0, sink_recache_after_switch=False): if self.model_type == 'i2v': assert clip_fea is not None and y is not None device = self.patch_embedding.weight.device if self.freqs.device != device: self.freqs = self.freqs.to(device) if y is not None: x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)] x = [self.patch_embedding(u.unsqueeze(0)) for u in x] grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) x = [u.flatten(2).transpose(1, 2) for u in x] seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) assert seq_lens.max() <= seq_len x = torch.cat(x) e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x)) e0 = self.time_projection(e).unflatten(1, (6, self.dim)).unflatten(dim=0, sizes=t.shape) context_lens = None context = self.text_embedding(torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])) if clip_fea is not None: context_clip = self.img_emb(clip_fea) context = torch.concat([context_clip, context], dim=1) if kv_cache is not None: first_global_end = kv_cache[0]['global_end_index'].item() if kv_cache else 0 if first_global_end == 0 and current_start == 0: B = x.shape[0] if self.query_memory_encoder is not None: self.query_memory_encoder.reset(batch_size=B, device=x.device, dtype=torch.bfloat16) self._ei_prev_window_start = None enc = self.query_memory_encoder num_blocks = len(self.blocks) if enc is not None and getattr(enc, 'has_history', False): num_query_groups = getattr(enc, 'num_query_groups', 1) if num_query_groups > 1: blocks_per_group = num_blocks // num_query_groups group_kvs = [enc.get_kv(group_index=g) for g in range(num_query_groups)] memory_kv_list = [group_kvs[min(i // blocks_per_group, num_query_groups - 1)] for i in range(num_blocks)] else: kv = enc.get_kv() memory_kv_list = [kv] * num_blocks else: memory_kv_list = [None] * num_blocks kwargs = dict(e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs, context=context, context_lens=context_lens, block_mask=self.block_mask, sink_recache_after_switch=sink_recache_after_switch, memory_kv=None) def create_custom_forward(module): def custom_forward(*inputs, **_kwargs): return module(*inputs, **_kwargs) return custom_forward cache_update_info = None cache_update_infos = [] for block_index, block in enumerate(self.blocks): kwargs['memory_kv'] = memory_kv_list[block_index] if torch.is_grad_enabled() and self.gradient_checkpointing: kwargs.update({'kv_cache': kv_cache[block_index], 'current_start': current_start, 'cache_start': cache_start}) result = torch.utils.checkpoint.checkpoint(create_custom_forward(block), x, **kwargs, use_reentrant=False) if kv_cache is not None and isinstance(result, tuple): x, block_cache_update_info = result cache_update_infos.append((block_index, block_cache_update_info)) cache_update_info = block_cache_update_info[:2] else: x = result else: kwargs.update({'kv_cache': kv_cache[block_index], 'crossattn_cache': crossattn_cache[block_index], 'current_start': current_start, 'cache_start': cache_start}) result = block(x, **kwargs) if kv_cache is not None and isinstance(result, tuple): x, block_cache_update_info = result cache_update_infos.append((block_index, block_cache_update_info)) cache_update_info = block_cache_update_info[:2] else: x = result if kv_cache is not None and cache_update_infos: if self.query_memory_encoder is not None: last_block_idx = cache_update_infos[-1][0] last_update_info = cache_update_infos[-1][1] update_dict = last_update_info[2] if len(last_update_info) > 2 else None sink_tok = self.blocks[0].self_attn.sink_size * 1560 sink_frames = self.blocks[0].self_attn.sink_size frame_seqlen = 1560 max_attn = self.blocks[0].self_attn.max_attention_size recent_window_frames = (max_attn - sink_tok) // frame_seqlen num_new_frames = (cache_update_infos[-1][1][0] - current_start) // frame_seqlen current_end_frame = current_start // frame_seqlen + num_new_frames oldest_recent_frame = max(sink_frames, current_end_frame - recent_window_frames) prev_oldest = self._ei_prev_window_start if self._ei_prev_window_start is not None else sink_frames num_exited_frames = max(0, oldest_recent_frame - prev_oldest) self._ei_prev_window_start = oldest_recent_frame if num_exited_frames > 0: num_exited_tokens = num_exited_frames * frame_seqlen if update_dict is not None and update_dict.get('action') == 'roll_and_insert': capture_start = sink_tok else: capture_start = prev_oldest * frame_seqlen cache_blk = kv_cache[last_block_idx] exited_k = cache_blk['k'][:, capture_start:capture_start + num_exited_tokens].clone() exited_v = cache_blk['v'][:, capture_start:capture_start + num_exited_tokens].clone() sink_k = cache_blk['k'][:, :sink_tok].clone() if sink_tok > 0 else None sink_v = cache_blk['v'][:, :sink_tok].clone() if sink_tok > 0 else None self.query_memory_encoder.update(exited_k, exited_v, sink_k, sink_v) self._apply_cache_updates(kv_cache, cache_update_infos) x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2)) x = self.unpatchify(x, grid_sizes) return torch.stack(x)