fix: add gate_compress parameter to all attention backends (#817)

fastvideo/attention/backends/abstract.py

@@ -102,6 +102,8 @@ def forward(
         value: torch.Tensor,
         kv_cache: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        gate_compress: torch.Tensor | None = None,
+
     ) -> torch.Tensor:
         ...
 
@@ -169,5 +171,6 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: T,
+        gate_compress: torch.Tensor | None = None,
     ) -> torch.Tensor:
         raise NotImplementedError

fastvideo/attention/backends/flash_attn.py

@@ -100,6 +100,8 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: FlashAttnMetadata,
+        gate_compress: torch.Tensor | None = None,
+
     ):
 
         def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor:

fastvideo/attention/backends/sage_attn.py

@@ -53,6 +53,7 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        gate_compress: torch.Tensor | None = None,
     ) -> torch.Tensor:
         output = sageattn(
             query,

fastvideo/attention/backends/sage_attn3.py

@@ -61,6 +61,8 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        gate_compress: torch.Tensor | None = None,
+
     ) -> torch.Tensor:
         query = query.transpose(1, 2)
         key = key.transpose(1, 2)

fastvideo/attention/backends/sdpa.py

@@ -74,6 +74,7 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: SDPAMetadata,
+        gate_compress: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # transpose to bs, heads, seq_len, head_dim
         query = query.transpose(1, 2)

fastvideo/attention/backends/sla.py

@@ -269,6 +269,7 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        gate_compress: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Forward pass for SLA attention.
 
@@ -463,6 +464,7 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        gate_compress: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Forward pass for SageSLA attention with quantized kernels.
 

fastvideo/attention/backends/video_sparse_attn.py

@@ -232,13 +232,13 @@ def postprocess_output(
     ) -> torch.Tensor:
         return self.untile(output, attn_metadata.reverse_tile_partition_indices, attn_metadata.non_pad_index)
 
-    def forward(  # type: ignore[override]
+    def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        gate_compress: torch.Tensor,
         attn_metadata: VideoSparseAttentionMetadata,
+        gate_compress: torch.Tensor,
     ) -> torch.Tensor:
         query = query.transpose(1, 2).contiguous()
         key = key.transpose(1, 2).contiguous()

fastvideo/attention/backends/vmoba.py

@@ -143,6 +143,8 @@ def forward(
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        gate_compress: torch.Tensor | None = None,
+
     ) -> torch.Tensor:
         """
         query: [B, L, H, D]