# Copyright 2024-2025 LongLive Authors. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ Causal Wan2.2-TI2V-5B model with Ulysses-style sequence parallelism. This inference-only variant mirrors the regular CausalWanModel parameter layout so existing generator checkpoints can be loaded into ``model.*`` keys. """ import math import os import torch import torch.nn as nn import torch.nn.functional as F from diffusers.configuration_utils import ConfigMixin, register_to_config from diffusers.models.modeling_utils import ModelMixin from wan_5b.modules.attention import attention from wan_5b.modules.model import ( WanRMSNorm, WanLayerNorm, rope_apply, rope_params, sinusoidal_embedding_1d, ) from wan_5b.modules.causal_model import causal_rope_apply, MultiShotT2VCrossAttention from wan_5b.distributed.sp_ulysses_inference import ( get_sp_rank, get_sp_world_size, is_sp_enabled, sp_all_gather, sp_scatter, ulysses_head_to_seq, ulysses_seq_to_head, ) QK_QUANT_SIM = os.environ.get("QK_QUANT_SIM", "0") == "1" QK_QUANT_RATIO = float(os.environ.get("QK_QUANT_RATIO", "0.28125")) try: import torch.cuda.nvtx as nvtx NVTX_ENABLED = True except Exception: nvtx = None NVTX_ENABLED = False class NVTXRange: def __init__(self, name: str): self.name = name def __enter__(self): if NVTX_ENABLED: nvtx.range_push(self.name) def __exit__(self, exc_type, exc, tb): if NVTX_ENABLED: nvtx.range_pop() def _get_d_quant(head_dim: int, ratio: float = QK_QUANT_RATIO) -> int: d_quant = max(8, round(head_dim * ratio)) return ((d_quant + 7) // 8) * 8 def _compute_ulysses_frame_info(num_frames, local_seq_len, sp_world_size, sp_rank): if num_frames == 1: return 1, local_seq_len, 0 assert num_frames % sp_world_size == 0, ( f"Ulysses SP requires num_frames ({num_frames}) divisible by " f"world_size ({sp_world_size})." ) local_frames = num_frames // sp_world_size frame_seqlen = local_seq_len // local_frames frame_offset = sp_rank * local_frames return local_frames, frame_seqlen, frame_offset class UlyssesCausalWanSelfAttention(nn.Module): """Causal self-attention with Ulysses sequence/head exchange.""" def __init__(self, dim, num_heads, local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6): assert dim % num_heads == 0 super().__init__() self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads self.local_attn_size = local_attn_size self.sink_size = sink_size self.global_sink_size = 0 self.qk_norm = qk_norm self.eps = eps if not isinstance(local_attn_size, int) and hasattr(local_attn_size, "__iter__"): values = list(local_attn_size) else: values = [int(local_attn_size)] non_neg_vals = [int(v) for v in values if int(v) != -1] max_local = max(non_neg_vals) if non_neg_vals else -1 self.max_attention_size = 28160 if max_local == -1 else max_local * 880 self.q = nn.Linear(dim, dim) self.k = nn.Linear(dim, dim) self.v = nn.Linear(dim, dim) self.o = nn.Linear(dim, dim) self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache=None, current_start=0, cache_start=None, t_scale=1.0, use_relative_rope=False, method="linear", original_seq_len=None, temporal_offset=0.0): b, s_local = x.shape[:2] n, d = self.num_heads, self.head_dim if cache_start is None: cache_start = current_start with NVTXRange("ulysses_qkv_proj"): q = self.norm_q(self.q(x)).view(b, s_local, n, d) k = self.norm_k(self.k(x)).view(b, s_local, n, d) v = self.v(x).view(b, s_local, n, d) if kv_cache is None: return self._forward_no_cache(q, k, v, grid_sizes, freqs, t_scale, method, original_seq_len) return self._forward_with_cache( q, k, v, grid_sizes, freqs, kv_cache, current_start, cache_start, t_scale, use_relative_rope, method, original_seq_len, temporal_offset, ) def _exchange_qkv(self, q, k, v): if QK_QUANT_SIM and is_sp_enabled(): d = q.shape[-1] d_quant = _get_d_quant(d) q_heads = F.pad(ulysses_seq_to_head(q[..., :d_quant]), (0, d - d_quant)) k_heads = F.pad(ulysses_seq_to_head(k[..., :d_quant]), (0, d - d_quant)) v_heads = ulysses_seq_to_head(v) else: q_heads = ulysses_seq_to_head(q) k_heads = ulysses_seq_to_head(k) v_heads = ulysses_seq_to_head(v) return q_heads, k_heads, v_heads def _forward_no_cache(self, q, k, v, grid_sizes, freqs, t_scale, method, original_seq_len): with NVTXRange("ulysses_seq_to_head"): q_heads, k_heads, v_heads = self._exchange_qkv(q, k, v) with NVTXRange("ulysses_rope"): roped_q = rope_apply( q_heads, grid_sizes, freqs, t_scale=t_scale, method=method, original_seq_len=original_seq_len, ).type_as(v_heads) roped_k = rope_apply( k_heads, grid_sizes, freqs, t_scale=t_scale, method=method, original_seq_len=original_seq_len, ).type_as(v_heads) with NVTXRange("ulysses_attention"): out_heads = attention(roped_q, roped_k, v_heads, causal=True) with NVTXRange("ulysses_head_to_seq"): out = ulysses_head_to_seq(out_heads) with NVTXRange("ulysses_out_proj"): return self.o(out.flatten(2)) def _forward_with_cache(self, q, k, v, grid_sizes, freqs, kv_cache, current_start, cache_start, t_scale, use_relative_rope, method, original_seq_len, temporal_offset): sp_world_size = get_sp_world_size() if is_sp_enabled() else 1 s_total = q.shape[1] * sp_world_size frame_seqlen = math.prod(grid_sizes[0][1:]).item() current_start_frame = current_start // frame_seqlen current_end = current_start + s_total with NVTXRange("ulysses_seq_to_head"): q_heads, k_heads, v_heads = self._exchange_qkv(q, k, v) with NVTXRange("ulysses_rope_cached"): roped_q = causal_rope_apply( q_heads, grid_sizes, freqs, start_frame=current_start_frame, t_scale=t_scale, method=method, original_seq_len=original_seq_len, temporal_offset=temporal_offset, ).type_as(v_heads) if use_relative_rope: key_to_cache = k_heads else: key_to_cache = causal_rope_apply( k_heads, grid_sizes, freqs, start_frame=current_start_frame, t_scale=t_scale, method=method, original_seq_len=original_seq_len, temporal_offset=temporal_offset, ).type_as(v_heads) with NVTXRange("ulysses_cache_update"): k_full, v_full = self._update_cache_and_get_kv( key_to_cache, v_heads, kv_cache, current_start, current_end, frame_seqlen, ) if use_relative_rope: raise NotImplementedError("use_relative_rope is not implemented for SP inference.") with NVTXRange("ulysses_cached_attention"): out_heads = attention(roped_q, k_full, v_full, causal=False) with NVTXRange("ulysses_head_to_seq"): out = ulysses_head_to_seq(out_heads) with NVTXRange("ulysses_out_proj"): out = self.o(out.flatten(2)) return out, (current_end, kv_cache["local_end_index"].item(), None) def _effective_sink(self, kv_cache, frame_seqlen): sink_tokens = self.sink_size * frame_seqlen global_sink_tokens = getattr(self, "global_sink_size", 0) * frame_seqlen pinned_start_t = kv_cache.get("pinned_start", None) if pinned_start_t is not None and hasattr(pinned_start_t, "item"): pinned_start = pinned_start_t.item() pinned_len = kv_cache["pinned_len"].item() else: pinned_start = -1 pinned_len = 0 has_pinned = pinned_start >= 0 and pinned_len > 0 if has_pinned and pinned_start == global_sink_tokens: effective_sink = global_sink_tokens + pinned_len elif has_pinned: effective_sink = global_sink_tokens else: effective_sink = max(global_sink_tokens, sink_tokens) return effective_sink, pinned_start, pinned_len, has_pinned def _update_cache_and_get_kv(self, k_new, v_new, kv_cache, current_start, current_end, frame_seqlen): b, s_new, _, _ = k_new.shape kv_cache_size = kv_cache["k"].shape[1] global_end_prev = kv_cache["global_end_index"].item() local_end_prev = kv_cache["local_end_index"].item() is_recompute = current_end <= global_end_prev and current_start > 0 effective_sink, pinned_start, pinned_len, has_pinned = self._effective_sink(kv_cache, frame_seqlen) need_roll = ( self.local_attn_size != -1 and current_end > global_end_prev and s_new + local_end_prev > kv_cache_size ) if need_roll: num_evicted = s_new + local_end_prev - kv_cache_size num_rolled = local_end_prev - num_evicted - effective_sink if num_rolled > 0: kv_cache["k"][:, effective_sink:effective_sink + num_rolled] = ( kv_cache["k"][ :, effective_sink + num_evicted:effective_sink + num_evicted + num_rolled ].clone() ) kv_cache["v"][:, effective_sink:effective_sink + num_rolled] = ( kv_cache["v"][ :, effective_sink + num_evicted:effective_sink + num_evicted + num_rolled ].clone() ) if has_pinned and pinned_start >= effective_sink and "pinned_start" in kv_cache: kv_cache["pinned_start"].sub_(num_evicted) pinned_start -= num_evicted local_end_new = local_end_prev + (current_end - global_end_prev) - num_evicted else: local_end_new = local_end_prev + (current_end - global_end_prev) local_start_new = local_end_new - s_new write_start = max(local_start_new, effective_sink) if is_recompute else local_start_new write_offset = max(0, write_start - local_start_new) write_len = max(0, local_end_new - write_start) if write_len > 0: kv_cache["k"][:, write_start:local_end_new] = k_new[:, write_offset:write_offset + write_len] kv_cache["v"][:, write_start:local_end_new] = v_new[:, write_offset:write_offset + write_len] if not is_recompute: kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_new) window_start = max(0, local_end_new - self.max_attention_size) prepend_sink = effective_sink > 0 and window_start > 0 prepend_pinned = has_pinned and pinned_start >= effective_sink and pinned_start < window_start if prepend_sink and prepend_pinned: extra = effective_sink + pinned_len effective_local_size = max(0, self.max_attention_size - extra) local_window_start = max(effective_sink, local_end_new - effective_local_size) k_full = torch.cat([ kv_cache["k"][:, :effective_sink], kv_cache["k"][:, pinned_start:pinned_start + pinned_len], kv_cache["k"][:, local_window_start:local_end_new], ], dim=1) v_full = torch.cat([ kv_cache["v"][:, :effective_sink], kv_cache["v"][:, pinned_start:pinned_start + pinned_len], kv_cache["v"][:, local_window_start:local_end_new], ], dim=1) elif prepend_sink: effective_local_size = max(0, self.max_attention_size - effective_sink) local_window_start = max(effective_sink, local_end_new - effective_local_size) k_full = torch.cat([ kv_cache["k"][:, :effective_sink], kv_cache["k"][:, local_window_start:local_end_new], ], dim=1) v_full = torch.cat([ kv_cache["v"][:, :effective_sink], kv_cache["v"][:, local_window_start:local_end_new], ], dim=1) elif prepend_pinned: effective_local_size = max(0, self.max_attention_size - pinned_len) local_window_start = max(0, local_end_new - effective_local_size) k_full = torch.cat([ kv_cache["k"][:, pinned_start:pinned_start + pinned_len], kv_cache["k"][:, local_window_start:local_end_new], ], dim=1) v_full = torch.cat([ kv_cache["v"][:, pinned_start:pinned_start + pinned_len], kv_cache["v"][:, local_window_start:local_end_new], ], dim=1) else: k_full = kv_cache["k"][:, window_start:local_end_new] v_full = kv_cache["v"][:, window_start:local_end_new] return k_full, v_full class UlyssesCausalWanAttentionBlock(nn.Module): def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, sink_size=0, qk_norm=True, cross_attn_norm=False, eps=1e-6): super().__init__() self.norm1 = WanLayerNorm(dim, eps) self.self_attn = UlyssesCausalWanSelfAttention( dim, num_heads, local_attn_size, sink_size, qk_norm, eps ) self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() self.cross_attn = MultiShotT2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps) self.norm2 = WanLayerNorm(dim, eps) self.ffn = nn.Sequential( nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim) ) self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens, kv_cache=None, crossattn_cache=None, current_start=0, cache_start=None, num_frames=None, local_frame_offset=0, t_scale=1.0, use_relative_rope=False, method="linear", original_seq_len=None, temporal_offset=0.0): num_frames_total = e.shape[1] if num_frames is None else num_frames local_seq_len = x.shape[1] if is_sp_enabled() and local_seq_len > 0: actual_local_frames, local_frame_seqlen, _ = _compute_ulysses_frame_info( num_frames_total, local_seq_len, get_sp_world_size(), get_sp_rank() ) e_local = e[:, local_frame_offset:local_frame_offset + actual_local_frames] else: actual_local_frames = num_frames_total local_frame_seqlen = local_seq_len // actual_local_frames if actual_local_frames > 0 else 0 e_local = e e_chunks = (self.modulation.unsqueeze(1) + e_local).chunk(6, dim=2) x_normed = self.norm1(x) if actual_local_frames > 0 and local_frame_seqlen > 0: x_mod = x_normed.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) x_mod = (x_mod * (1 + e_chunks[1]) + e_chunks[0]).flatten(1, 2) else: x_mod = x_normed self_attn_result = self.self_attn( x_mod, seq_lens, grid_sizes, freqs, kv_cache, current_start, cache_start, t_scale, use_relative_rope, method, original_seq_len, temporal_offset, ) if kv_cache is not None and isinstance(self_attn_result, tuple): y, cache_update_info = self_attn_result else: y = self_attn_result cache_update_info = None if actual_local_frames > 0 and local_frame_seqlen > 0: y_mod = y.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) x = x + (y_mod * e_chunks[2]).flatten(1, 2) else: x = x + y if x.shape[1] > 0: x = x + self.cross_attn( self.norm3(x), context, context_lens, crossattn_cache=crossattn_cache ) x_normed = self.norm2(x) if actual_local_frames > 0 and local_frame_seqlen > 0: x_mod = x_normed.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) y = self.ffn((x_mod * (1 + e_chunks[4]) + e_chunks[3]).flatten(1, 2)) y_mod = y.unflatten(dim=1, sizes=(actual_local_frames, local_frame_seqlen)) x = x + (y_mod * e_chunks[5]).flatten(1, 2) else: x = x + self.ffn(x_normed) return (x, cache_update_info) if cache_update_info is not None else x class UlyssesCausalHead(nn.Module): def __init__(self, dim, out_dim, patch_size, eps=1e-6): super().__init__() self.dim = dim self.out_dim = out_dim self.patch_size = patch_size self.norm = WanLayerNorm(dim, eps) self.head = nn.Linear(dim, math.prod(patch_size) * out_dim) self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) def forward(self, x, e, local_frame_offset=0, actual_local_frames=None, frame_seqlen_global=None): num_frames_total = e.shape[1] local_seq_len = x.shape[1] if is_sp_enabled() and actual_local_frames is not None: local_frames = actual_local_frames local_seqlen = local_seq_len // local_frames if local_frames > 0 else frame_seqlen_global e_local = e[:, local_frame_offset:local_frame_offset + local_frames] else: local_frames = num_frames_total local_seqlen = local_seq_len // local_frames if local_frames > 0 else 0 e_local = e if local_frames == 0 or local_seq_len == 0: b, _, _ = x.shape return torch.empty( b, 0, local_seqlen if local_seqlen else 1, math.prod(self.patch_size) * self.out_dim, dtype=x.dtype, device=x.device, ) e_chunks = (self.modulation.unsqueeze(1) + e_local).chunk(2, dim=2) x = self.norm(x).unflatten(dim=1, sizes=(local_frames, local_seqlen)) return self.head(x * (1 + e_chunks[1]) + e_chunks[0]) class UlyssesSPCausalWanModel(ModelMixin, ConfigMixin): ignore_for_config = ["patch_size", "cross_attn_norm", "qk_norm", "text_dim"] _no_split_modules = ["UlyssesCausalWanAttentionBlock"] _supports_gradient_checkpointing = True @register_to_config def __init__(self, model_type="ti2v", patch_size=(1, 2, 2), text_len=512, in_dim=48, dim=3072, ffn_dim=14336, freq_dim=256, text_dim=4096, out_dim=48, num_heads=24, num_layers=30, local_attn_size=-1, sink_size=0, num_frame_per_block=1, qk_norm=True, cross_attn_norm=True, eps=1e-6): super().__init__() assert model_type in ["t2v", "i2v", "ti2v"] self.model_type = model_type self.patch_size = patch_size self.text_len = text_len self.in_dim = in_dim self.dim = dim self.ffn_dim = ffn_dim self.freq_dim = freq_dim self.text_dim = text_dim self.out_dim = out_dim self.num_heads = num_heads self.num_layers = num_layers self.local_attn_size = local_attn_size self.sink_size = sink_size self.qk_norm = qk_norm self.cross_attn_norm = cross_attn_norm self.eps = eps self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) self.text_embedding = nn.Sequential( nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) ) self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) self.blocks = nn.ModuleList([ UlyssesCausalWanAttentionBlock( dim, ffn_dim, num_heads, local_attn_size, sink_size, qk_norm, cross_attn_norm, eps, ) for _ in range(num_layers) ]) self.head = UlyssesCausalHead(dim, out_dim, patch_size, eps) d = dim // num_heads self.freqs = torch.cat([ rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6)), ], dim=1) self.init_weights() self.gradient_checkpointing = False self.num_frame_per_block = num_frame_per_block self.t_scale = 1.0 self.use_relative_rope = False self.rope_method = "linear" self.original_seq_len = None self.rope_temporal_offset = 0.0 self.kv_quant_config = None def forward(self, x, t, context, seq_len, clip_fea=None, y=None, kv_cache=None, crossattn_cache=None, current_start=0, cache_start=0): sp_world_size = get_sp_world_size() if is_sp_enabled() else 1 sp_rank = get_sp_rank() if is_sp_enabled() else 0 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)] with NVTXRange("ulysses_patch_embedding"): 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) num_frames = t.shape[1] if t.dim() > 1 else grid_sizes[0, 0].item() frame_seqlen = x.shape[1] // num_frames if is_sp_enabled(): assert self.num_heads % sp_world_size == 0 assert x.shape[1] % sp_world_size == 0 with NVTXRange("ulysses_scatter_input"): x = sp_scatter(x, dim=1) actual_local_frames, local_frame_seqlen, local_frame_offset = _compute_ulysses_frame_info( num_frames, x.shape[1], sp_world_size, sp_rank ) else: actual_local_frames = num_frames local_frame_seqlen = frame_seqlen local_frame_offset = 0 with NVTXRange("ulysses_time_embedding"): 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) with NVTXRange("ulysses_text_embedding"): 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 = torch.concat([self.img_emb(clip_fea), context], dim=1) kwargs = dict( e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs, context=context, context_lens=context_lens, num_frames=num_frames, local_frame_offset=local_frame_offset, t_scale=self.t_scale, use_relative_rope=self.use_relative_rope, method=self.rope_method, original_seq_len=self.original_seq_len, temporal_offset=self.rope_temporal_offset, ) for block_idx, block in enumerate(self.blocks): kwargs.update({ "kv_cache": kv_cache[block_idx] if kv_cache else None, "crossattn_cache": crossattn_cache[block_idx] if crossattn_cache else None, "current_start": current_start, "cache_start": cache_start, }) result = block(x, **kwargs) x = result[0] if kv_cache is not None and isinstance(result, tuple) else result x = self.head( x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2), local_frame_offset=local_frame_offset, actual_local_frames=actual_local_frames, frame_seqlen_global=frame_seqlen, ) if is_sp_enabled(): x = sp_all_gather(x, dim=2 if num_frames == 1 else 1) x = self.unpatchify(x, grid_sizes) return torch.stack(x) def unpatchify(self, x, grid_sizes): c = self.out_dim pt, ph, pw = self.patch_size output = [] for i, (f, h, w) in enumerate(grid_sizes.tolist()): x_i = x[i, :f, :h * w, :].reshape(f, h, w, pt, ph, pw, c) x_i = x_i.permute(6, 0, 3, 1, 4, 2, 5).reshape(c, f * pt, h * ph, w * pw) output.append(x_i) return output def init_weights(self): def _init(module): if isinstance(module, nn.Linear): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Conv3d): nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) self.apply(_init)