diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index 9d716e9bc..1a0c91668 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -20,7 +20,7 @@ def main(): ) sampling_param = SamplingParam.from_pretrained(model_name) - sampling_param.num_frames = 77 + sampling_param.num_frames = 13 prompt = ( "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " "wide with interest. The playful yet serene atmosphere is complemented by soft " diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 2046f99d4..4ca94e31e 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -67,7 +67,10 @@ class DistributedAttention(nn.Module): replicated_q: torch.Tensor | None = None, replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + k_cache: torch.Tensor | None = None, + v_cache: torch.Tensor | None = None, + return_current_kv: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """Forward pass for distributed attention. Args: @@ -103,11 +106,23 @@ class DistributedAttention(nn.Module): qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1) + # q = sequence_model_parallel_all_to_all_4D(q, + # scatter_dim=2, + # gather_dim=1) + # k = sequence_model_parallel_all_to_all_4D(k, + # scatter_dim=2, + # gather_dim=1) + # v = sequence_model_parallel_all_to_all_4D(v, + # scatter_dim=2, + # gather_dim=1) + + # Apply backend-specific preprocess_qkv qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) # Concatenate with replicated QKV if provided if replicated_q is not None: + assert False, "should not be used" assert replicated_k is not None and replicated_v is not None replicated_qkv = torch.cat( [replicated_q, replicated_k, replicated_v], @@ -118,9 +133,30 @@ class DistributedAttention(nn.Module): heads_per_rank] qkv = torch.cat([qkv, replicated_qkv], dim=1) - q, k, v = qkv.chunk(3, dim=0) + q, k_new, v_new = qkv.chunk(3, dim=0) + logger.info( + f"rank: {local_rank}, k_new: {k_new.shape}, v_new: {v_new.shape}", + local_main_process_only=False) + k_total = k_new + v_total = v_new - output = self.attn_impl.forward(q, k, v, ctx_attn_metadata) + if k_cache is not None or v_cache is not None: + assert False, "KV cache is not supported" + assert k_cache is not None and v_cache is not None + assert k_cache.shape == v_cache.shape + logger.info( + f"rank: {local_rank}, k_cache: {k_cache.shape}, v_cache: {v_cache.shape}", + local_main_process_only=False) + assert k_cache.shape[2] == v_cache.shape[2], "Number of heads must be the same" + if k_cache.shape[1] > 0: + k_total = torch.cat([k_cache, k_new], dim=1) + v_total = torch.cat([v_cache, v_new], dim=1) + assert k_total.shape == v_total.shape, "Key and value shapes must be the same" + + logger.info( + f"rank: {local_rank}, k_total: {k_total.shape}, v_total: {v_total.shape}", + local_main_process_only=False) + output = self.attn_impl.forward(q, k_total, v_total, ctx_attn_metadata) # Redistribute back if using sequence parallelism replicated_output = None @@ -136,7 +172,10 @@ class DistributedAttention(nn.Module): output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2) - return output, replicated_output + if return_current_kv: + return output, replicated_output, k_new, v_new # [batch_size, seq_len, num_heads/sp_world_size, head_dim] + else: + return output, replicated_output, None, None class DistributedAttention_VSA(DistributedAttention): @@ -153,7 +192,8 @@ class DistributedAttention_VSA(DistributedAttention): replicated_k: torch.Tensor | None = None, replicated_v: torch.Tensor | None = None, gate_compress: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None]: + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: """Forward pass for distributed attention. Args: @@ -170,8 +210,15 @@ class DistributedAttention_VSA(DistributedAttention): - o (torch.Tensor): Output tensor after attention for the main sequence - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided """ - # Check text tokens are not supported for VSA now + k_cache = kwargs.pop("k_cache", None) + v_cache = kwargs.pop("v_cache", None) + return_current_kv = kwargs.pop("return_current_kv", False) + assert not kwargs, f"Unexpected kwargs for DistributedAttention_VSA: {kwargs.keys()}" + + # Check unsupported features assert replicated_q is None and replicated_k is None and replicated_v is None, "Replicated QKV is not supported for VSA now" + assert not return_current_kv, "return_current_kv is not supported for VSA attention" + assert k_cache is None and v_cache is None, "KV cache is not supported for VSA attention" # Check input shapes assert q.dim() == 4 and k.dim() == 4 and v.dim( ) == 4, "Expected 4D tensors" @@ -203,7 +250,7 @@ class DistributedAttention_VSA(DistributedAttention): output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2) - return output, replicated_output + return output, replicated_output, None, None class LocalAttention(nn.Module): diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 69b3e5430..fac623f80 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -20,6 +20,7 @@ import fastvideo.envs as envs from fastvideo.attention import (DistributedAttention, LocalAttention) from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.distributed import sequence_model_parallel_all_gather from fastvideo.distributed.parallel_state import get_sp_world_size from fastvideo.forward_context import get_forward_context from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, @@ -59,7 +60,7 @@ class CausalWanSelfAttention(nn.Module): self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560 # Scaled dot product attention - self.attn = LocalAttention( + self.attn = DistributedAttention( num_heads=num_heads, head_size=self.head_dim, dropout_rate=0, @@ -92,6 +93,7 @@ class CausalWanSelfAttention(nn.Module): roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) if kv_cache is None: + assert False, "should not be used" # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] padded_roped_query = torch.cat( @@ -121,44 +123,82 @@ class CausalWanSelfAttention(nn.Module): )[:, :, :-padded_length].transpose(2, 1) else: frame_seqlen = q.shape[1] - current_end = current_start + roped_query.shape[1] sink_tokens = self.sink_size * frame_seqlen - # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache kv_cache_size = kv_cache["k"].shape[1] - num_new_tokens = roped_query.shape[1] - 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): - # Calculate the number of new tokens added in this step - # Shift existing cache content left to discard oldest tokens - # Clone the source slice to avoid overlapping memory error - 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 - kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ - kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() - # Insert the new keys/values at the end - 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 - kv_cache["k"][:, local_start_index:local_end_index] = roped_key - kv_cache["v"][:, local_start_index:local_end_index] = v + sp_world_size = max(get_sp_world_size(), 1) + num_new_tokens = roped_query.shape[1] * sp_world_size + from fastvideo.distributed.parallel_state import get_sp_group + sp_group = get_sp_group() + sp_rank = sp_group.rank_in_group + logger.info('rank: %s, num_new_tokens: %s', sp_rank, num_new_tokens, local_main_process_only=False) + + prev_global_end = kv_cache["global_end_index"].item() + # prev_local_end = kv_cache["local_end_index"].item() + prev_local_end = current_start + # prev_local_end = + # logger.info('prev_global_end: %s, prev_local_end: %s', prev_global_end, prev_local_end) + + current_end = current_start + num_new_tokens + logger.info('rank: %s, current_end: %s', sp_rank, current_end, local_main_process_only=False) + + assert self.local_attn_size == -1, "Local attention is not supported" + # if self.local_attn_size != -1 and (current_end > prev_global_end) and ( + # num_new_tokens + prev_local_end > kv_cache_size): + # num_evicted_tokens = num_new_tokens + prev_local_end - kv_cache_size + # num_rolled_tokens = prev_local_end - num_evicted_tokens - sink_tokens + # if num_rolled_tokens > 0: + # src_slice = slice(sink_tokens + num_evicted_tokens, + # sink_tokens + num_evicted_tokens + num_rolled_tokens) + # dst_slice = slice(sink_tokens, sink_tokens + num_rolled_tokens) + # kv_cache["k"][:, dst_slice] = kv_cache["k"][:, src_slice].clone() + # kv_cache["v"][:, dst_slice] = kv_cache["v"][:, src_slice].clone() + # prev_local_end = max(sink_tokens, prev_local_end - num_evicted_tokens) + + # current_end = current_start + roped_query.shape[1] + 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 + max_history_tokens = max(self.max_attention_size - num_new_tokens, 0) + history_len = min(prev_local_end, max_history_tokens) + history_start = prev_local_end - history_len + logger.info('rank: %s, local_start_index: %s, local_end_index: %s', sp_rank, local_start_index, local_end_index, local_main_process_only=False) + logger.info('rank: %s, history_len: %s', sp_rank, history_len, local_main_process_only=False) + logger.info('rank: %s, history_start: %s', sp_rank, history_start, local_main_process_only=False) + if history_len > 0: + k_history = kv_cache["k"][:, history_start:history_start + history_len] + v_history = kv_cache["v"][:, history_start:history_start + history_len] + logger.info('k_history: %s, v_history: %s', k_history.shape, v_history.shape, local_main_process_only=False) else: - # Assign new keys/values directly up to current_end - 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 - kv_cache["k"] = kv_cache["k"].detach() - kv_cache["v"] = kv_cache["v"].detach() - # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None) - kv_cache["k"][:, local_start_index:local_end_index] = roped_key - kv_cache["v"][:, local_start_index:local_end_index] = v - x = self.attn( + k_history = None + v_history = None + + attn_output, _, k_new, v_new = self.attn( roped_query, - kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index], - kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index] + roped_key, + v, + k_cache=k_history, + v_cache=v_history, + return_current_kv=True, ) + + assert k_new.shape[1] == num_new_tokens == v_new.shape[1], "New keys and values must have the same number of tokens" + # assert actual_new_tokens == num_new_tokens, "Actual new tokens must match expected new tokens" + # num_new_tokens = actual_new_tokens + # current_end = current_start + num_new_tokens + + new_k = k_new.detach().to(kv_cache["k"].dtype) + new_v = v_new.detach().to(kv_cache["v"].dtype) + + + kv_cache["k"] = kv_cache["k"].detach() + kv_cache["v"] = kv_cache["v"].detach() + kv_cache["k"][:, local_start_index:local_end_index] = new_k + kv_cache["v"][:, local_start_index:local_end_index] = new_v + logger.info('rank: %s, local_start_index: %s, local_end_index: %s', sp_rank, local_start_index, local_end_index, local_main_process_only=False) + + x = attn_output kv_cache["global_end_index"].fill_(current_end) kv_cache["local_end_index"].fill_(local_end_index) + logger.info('FINAL rank: %s, global_end_index: %s, local_end_index: %s', sp_rank, kv_cache["global_end_index"], kv_cache["local_end_index"], local_main_process_only=False) return x diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 5a6ed36e4..2907f6912 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -248,6 +248,7 @@ class CausalDMDDenosingStage(DenoisingStage): current_latents = rearrange(current_latents, "b c (n t) h w -> b c n t h w", n=sp_world_size).contiguous() + logger.info(f"rearranged latents shape: {current_latents.shape}") current_latents = current_latents[:, :, rank_in_sp_group, :, :, :] logger.info( @@ -257,8 +258,10 @@ class CausalDMDDenosingStage(DenoisingStage): frames_per_rank = ( current_num_frames // sp_world_size) if sp_group else current_num_frames + # logger.info(f"frames_per_rank: {frames_per_rank}") rank_offset_frames = (rank_in_sp_group * frames_per_rank) if sp_group else 0 + # logger.info(f"rank {rank_in_sp_group} rank_offset_frames: {rank_offset_frames}", local_main_process_only=False) # use BTCHW for DMD conversion routines noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) video_raw_latent_shape = noise_latents_btchw.shape @@ -320,16 +323,16 @@ class CausalDMDDenosingStage(DenoisingStage): (latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long) + assert current_model is not None, "current_model is not initialized" pred_noise_btchw = current_model( latent_model_input, prompt_embeds, t_expanded_noise, kv_cache=_get_kv_cache(t_cur), crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index + rank_offset_frames) * + current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=pos_start_base + start_index + - rank_offset_frames, + start_frame=start_index, **image_kwargs, **pos_cond_kwargs, ).permute(0, 2, 1, 3, 4) @@ -415,7 +418,10 @@ class CausalDMDDenosingStage(DenoisingStage): t_context = torch.ones([latents.shape[0]], device=latents.device, dtype=torch.long) * int(context_noise) - # context_bcthw = current_latents.to(target_dtype) + context_bcthw = current_latents.to(target_dtype) + # context_frame_start = (pos_start_base + start_index) + # context_token_start = context_frame_start * self.frame_seq_length + with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \ @@ -425,6 +431,7 @@ class CausalDMDDenosingStage(DenoisingStage): t_expanded_context = t_context.unsqueeze(1) if boundary_timestep is not None: + assert self.transformer_2 is not None, "transformer_2 is not initialized" self.transformer_2( context_bcthw, prompt_embeds, @@ -444,10 +451,9 @@ class CausalDMDDenosingStage(DenoisingStage): t_expanded_context, kv_cache=kv_cache1, crossattn_cache=crossattn_cache, - current_start=(pos_start_base + start_index + rank_offset_frames) * + current_start=(pos_start_base + start_index) * self.frame_seq_length, - start_frame=pos_start_base + start_index + - rank_offset_frames, + start_frame=start_index, **image_kwargs, **pos_cond_kwargs, ) @@ -467,6 +473,9 @@ class CausalDMDDenosingStage(DenoisingStage): """ kv_cache1 = [] num_attention_heads = self.transformer.num_attention_heads + sp_world_size = get_sp_world_size() + assert num_attention_heads % max(sp_world_size, 1) == 0, "num_attention_heads must be divisible by sp_world_size" + heads_per_rank = num_attention_heads // max(sp_world_size, 1) attention_head_dim = self.transformer.attention_head_dim if self.local_attn_size != -1: kv_cache_size = self.local_attn_size * self.frame_seq_length @@ -477,14 +486,14 @@ class CausalDMDDenosingStage(DenoisingStage): kv_cache1.append({ "k": torch.zeros([ - batch_size, kv_cache_size, num_attention_heads, + batch_size, kv_cache_size, heads_per_rank, attention_head_dim ], dtype=dtype, device=device), "v": torch.zeros([ - batch_size, kv_cache_size, num_attention_heads, + batch_size, kv_cache_size, heads_per_rank, attention_head_dim ], dtype=dtype,