wip

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 "

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):

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
 

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,