[FA] Refactor FA CB kwargs

src/transformers/generation/continuous_batching/cache.py

@@ -11,9 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import inspect
 from math import floor, gcd, sqrt
-from typing import Any
 
 import torch
 
@@ -392,22 +390,6 @@ def update(
         # Return the new KV values
         return key_states_with_cache, value_states_with_cache
 
-    def get_block_table_key(self, flash_attn_with_kvcache_fn: Any) -> str:
-        """A function to get the name of the block table key for the given flash_attn_with_kvcache_fn. The function's
-        signature is only inspected once. This is necessary because different version of flash have different names for
-        the block table key."""
-        if self._block_table_key is None:
-            kwarg_names = inspect.signature(flash_attn_with_kvcache_fn).parameters.keys()
-            if "block_table" in kwarg_names:
-                self._block_table_key = "block_table"
-            elif "page_table" in kwarg_names:
-                self._block_table_key = "page_table"
-            else:
-                raise ValueError(
-                    f"flash_attn_with_kvcache_fn does not have a block_table or page_table argument: {inspect.signature(flash_attn_with_kvcache_fn)}"
-                )
-        return self._block_table_key
-
     def search_prefix_match(self, request_id: str, prompt_ids: list[int]) -> int:
         """Searches for a prefix match in the cache for the given (prompts_ids). If one is found, we reference the
         matching blocks in the (request_id), increase the reference count of the blocks and return the number of blocks

src/transformers/integrations/flash_attention.py

@@ -1,13 +1,11 @@
 import torch
 
-from ..modeling_flash_attention_utils import _flash_attention_forward, flash_attn_supports_top_left_mask
+from ..modeling_flash_attention_utils import _flash_attention_forward
 from ..utils import logging
 
 
 logger = logging.get_logger(__name__)
 
-_use_top_left_mask = flash_attn_supports_top_left_mask()
-
 
 def get_target_dtype(query: torch.Tensor, module: torch.nn.Module) -> torch.dtype:
     """If the query is in float32, return a target dtype compatible with flash attention. Return None otherwise."""
@@ -76,7 +74,6 @@ def flash_attention_forward(
         softmax_scale=scaling,
         sliding_window=sliding_window,
         softcap=softcap,
-        use_top_left_mask=_use_top_left_mask,
         target_dtype=target_dtype,
         attn_implementation=module.config._attn_implementation,
         layer_idx=module.layer_idx if hasattr(module, "layer_idx") else None,

src/transformers/integrations/flash_paged.py

@@ -4,6 +4,103 @@
 from ..modeling_flash_attention_utils import lazy_import_paged_flash_attention
 
 
+def _prepare_varlen_fa(
+    module: torch.nn.Module,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: PagedAttentionCache,
+    cu_seq_lens_q: torch.Tensor,
+    cu_seq_lens_k: torch.Tensor,
+    **kwargs,
+):
+    flash_attn_varlen_func, _, process_flash_kwargs_fn = lazy_import_paged_flash_attention(
+        module.config._attn_implementation
+    )
+
+    # .update changes the shape of k and v from [1, num_kv_heads, seqlen_kv, head_dim] to [seqlen_kv, num_kv_heads, head_dim]
+    k, v = cache.update(
+        key_states=k,
+        value_states=v,
+        layer_idx=module.layer_idx,
+        read_index=kwargs["read_index"],
+        write_index=kwargs["write_index"],
+    )
+    k = k.contiguous()
+    v = v.contiguous()
+
+    # [1, num_q_heads, seqlen_q, head_dim] to [seqlen_q, num_q_heads, head_dim]
+    q = q.transpose(1, 2).squeeze(0).contiguous()
+
+    # Prepare relevant args and kwargs
+    flash_args = (
+        q,
+        k,
+        v,
+    )
+
+    flash_kwargs = process_flash_kwargs_fn(
+        flash_fn_name="flash_varlen_fn",
+        **kwargs,
+    )
+    flash_kwargs = flash_kwargs | {
+        "cu_seqlens_q": cu_seq_lens_q.to(torch.int32),
+        "cu_seqlens_k": cu_seq_lens_k.to(torch.int32),
+    }
+
+    return flash_attn_varlen_func, flash_args, flash_kwargs
+
+
+def _prepare_kvcache_fa(
+    module: torch.nn.Module,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cache: PagedAttentionCache,
+    cu_seq_lens_k: torch.Tensor,
+    block_table: torch.Tensor | None,
+    **kwargs,
+):
+    # Retrieve the flash attention functions
+    _, flash_attn_with_kvcache, process_flash_kwargs_fn = lazy_import_paged_flash_attention(
+        module.config._attn_implementation
+    )
+
+    # Get layer group index for this layer
+    group_idx, layer_idx_in_group = cache.layer_index_to_group_indices[module.layer_idx]
+    # KV cache shape: [num_pages, num_kv_heads, head_dim] -> [num_blocks, block_size, num_kv_heads, head_dim]
+    k_cache = cache.key_cache[layer_idx_in_group].view(-1, cache.block_size, cache.num_key_value_heads, cache.head_dim)
+    v_cache = cache.value_cache[layer_idx_in_group].view(
+        -1, cache.block_size, cache.num_key_value_heads, cache.head_dim
+    )
+    # Reshape Q, K, V from [1, num_*_heads, batch_size, head_dim] to [batch_size, 1, num_*_heads, head_dim]
+    q = q.permute(2, 0, 1, 3).contiguous()
+    k = k.permute(2, 0, 1, 3).contiguous()
+    v = v.permute(2, 0, 1, 3).contiguous()
+    # Compute cache_seqlens from cu_seq_lens_k (current cache length BEFORE adding new tokens)
+    # cu_seq_lens_k is cumulative, so seqlens[i] = cu_seq_lens_k[i+1] - cu_seq_lens_k[i] - 1 (subtract 1 for the new token)
+    batch_size = k.size(0)
+    cache_seq_lens = cu_seq_lens_k[1 : batch_size + 1] - cu_seq_lens_k[:batch_size] - 1
+
+    # Prepare relevant args and kwargs
+    flash_args = (
+        q,
+        k_cache,
+        v_cache,
+        k,
+        v,
+    )
+
+    flash_kwargs = process_flash_kwargs_fn(
+        block_table=block_table[group_idx],
+        flash_fn_name="flash_with_kvcache_fn",
+        **kwargs,
+    )
+    flash_kwargs = flash_kwargs | {"cache_seqlens": cache_seq_lens.to(torch.int32)}
+
+    return flash_attn_with_kvcache, flash_args, flash_kwargs
+
+
 def paged_attention_forward(
     module: torch.nn.Module,
     q: torch.Tensor,
@@ -16,6 +113,8 @@ def paged_attention_forward(
     max_seqlen_q: int,
     max_seqlen_k: int | dict[str, int],
     block_table: torch.Tensor | None,
+    sliding_window: int | None = None,
+    is_causal: bool | None = None,
     **kwargs,
 ) -> tuple[torch.Tensor, None]:
     """Performs the forward pass of attention with paged key-value cache. This function handles the cache updates and
@@ -39,83 +138,56 @@ def paged_attention_forward(
             read and dispatches the read using the block table. Same for the write. If a request has less  blocksthan
             max_blocks_per_seq blocks, the block table is padded with -1s to indicate that this cache is not allocated.
     """
-    # Retrieve the flash attention functions
-    flash_attn_varlen_func, flash_attn_with_kvcache = lazy_import_paged_flash_attention(
-        module.config._attn_implementation
-    )
+    # Instead of relying on the value set in the module directly, we use the is_causal passed in kwargs if it is presented
+    is_causal = is_causal if is_causal is not None else module.is_causal
 
     # Retrieve the cumulative sequence lengths for the current layer
-    sliding_window = (-1, -1) if not getattr(module, "sliding_window", False) else (module.sliding_window - 1, 0)
-    layer_type = "full_attention" if sliding_window == (-1, -1) else "sliding_attention"
+    layer_type = "full_attention" if sliding_window is None else "sliding_attention"
     if isinstance(cu_seq_lens_k, dict):
         cu_seq_lens_k = cu_seq_lens_k[layer_type]
         max_seqlen_k = max_seqlen_k[layer_type]
 
     # If no block table is provided, use flash_attn_varlen_func with read/write indices
     if block_table is None:
-        # .update changes the shape of k and v from [1, num_kv_heads, seqlen_kv, head_dim] to [-1, num_kv_heads, head_dim]
-        k, v = cache.update(
-            key_states=k,
-            value_states=v,
-            layer_idx=module.layer_idx,
-            read_index=kwargs["read_index"],
-            write_index=kwargs["write_index"],
-        )
-        custom_kwargs = {"s_aux": kwargs.get("s_aux")} if "s_aux" in kwargs else {}
-        attn_output = flash_attn_varlen_func(
-            q.transpose(1, 2).squeeze(0).contiguous(),
-            k.contiguous(),
-            v.contiguous(),
-            cu_seq_lens_q.to(torch.int32),
-            cu_seq_lens_k.to(torch.int32).clone(),
-            max_seqlen_q,
-            max_seqlen_k,
-            softmax_scale=module.scaling,
-            causal=True,  # kind of a must, it automatically aligns the mask for q < k
-            window_size=sliding_window,  # -1 means infinite context window
-            **custom_kwargs,
+        fa_function, fa_args, fa_kwargs = _prepare_varlen_fa(
+            module,
+            q,
+            k,
+            v,
+            cache=cache,
+            cu_seq_lens_q=cu_seq_lens_q,
+            cu_seq_lens_k=cu_seq_lens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            sliding_window=sliding_window,
+            is_causal=is_causal,
+            **kwargs,
         )
-        if isinstance(attn_output, tuple):
-            attn_output = attn_output[0]
 
     # Otherwise, use flash_attn_with_kvcache which updates the cache in-place and computes attention
     else:
-        # Get layer group index for this layer
-        group_idx, layer_idx_in_group = cache.layer_index_to_group_indices[module.layer_idx]
-        # KV cache shape: [num_pages, num_kv_heads, head_dim] -> [num_blocks, block_size, num_kv_heads, head_dim]
-        k_cache = cache.key_cache[layer_idx_in_group].view(
-            -1, cache.block_size, cache.num_key_value_heads, cache.head_dim
-        )
-        v_cache = cache.value_cache[layer_idx_in_group].view(
-            -1, cache.block_size, cache.num_key_value_heads, cache.head_dim
-        )
-        # Reshape Q, K, V from [1, num_*_heads, batch_size, head_dim] to [batch_size, 1, num_*_heads, head_dim]
-        q = q.permute(2, 0, 1, 3).contiguous()
-        k = k.permute(2, 0, 1, 3).contiguous()
-        v = v.permute(2, 0, 1, 3).contiguous()
-        # Compute cache_seqlens from cu_seq_lens_k (current cache length BEFORE adding new tokens)
-        # cu_seq_lens_k is cumulative, so seqlens[i] = cu_seq_lens_k[i+1] - cu_seq_lens_k[i] - 1 (subtract 1 for the new token)
-        batch_size = k.size(0)
-        cache_seqlens = (cu_seq_lens_k[1 : batch_size + 1] - cu_seq_lens_k[:batch_size] - 1).to(torch.int32)
-        # The arg name for the block table is not the same in VLLM's kernel and Tri Dao's kernel, so we need to parse it
-        flash_kwargs = {cache.get_block_table_key(flash_attn_with_kvcache): block_table[group_idx]}
-        if "s_aux" in kwargs:
-            flash_kwargs["s_aux"] = kwargs["s_aux"]  # this is only available in VLLM's FA3
-        # Call flash_attn_with_kvcache - this updates cache in-place and computes attention
-        attn_output = flash_attn_with_kvcache(  # TODO: add more doc in a dedicated wrapper (coming in next PRs)
+        fa_function, fa_args, fa_kwargs = _prepare_kvcache_fa(
+            module,
             q,
-            k_cache,
-            v_cache,
-            k=k,
-            v=v,
-            cache_seqlens=cache_seqlens,
-            softmax_scale=module.scaling,
-            causal=True,
-            window_size=sliding_window,
-            **flash_kwargs,
+            k,
+            v,
+            cache=cache,
+            cu_seq_lens_k=cu_seq_lens_k,
+            block_table=block_table,
+            sliding_window=sliding_window,
+            is_causal=is_causal,
+            **kwargs,
         )
-        if isinstance(attn_output, tuple):
-            attn_output = attn_output[0]
+
+    # Run relevant underlying FA function
+    attn_output = fa_function(*fa_args, **fa_kwargs)
+
+    # Workaround for some kernels always returning tuples, e.g. vllm fa3 one
+    if isinstance(attn_output, tuple):
+        attn_output = attn_output[0]
+
+    if block_table is not None:
         # Reshape output from [batch_size, 1, num_heads, head_dim] to [batch_size, num_heads, head_dim]
         attn_output = attn_output.squeeze(1)
+
     return attn_output, None

src/transformers/integrations/npu_flash_attention.py

@@ -23,15 +23,19 @@
 
 
 # FlashAttention2 is supported on Ascend NPU with down-right aligned causal mask by default.
-# Set environment variable `NPU_FA2_SPARSE_MODE` to 2 when using top-left aligned causal mask.
 TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE = 2
 DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE = 3
 
 SPARSE_MODE = int(os.getenv("NPU_FA2_SPARSE_MODE", default=DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE))
-if SPARSE_MODE not in [TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE, DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE]:
+if SPARSE_MODE != DOWN_RIGHT_ALIGNED_CAUSAL_MASK_MODE:
+    if SPARSE_MODE == TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE:
+        raise ValueError(
+            "Environment variable `NPU_FA2_SPARSE_MODE` can only be set as 3 (down-right aligned causal mask). "
+            "Sparse mode 2 (top-left aligned causal mask) has been deprecated."
+        )
+
     raise ValueError(
-        "Environment variable `NPU_FA2_SPARSE_MODE` can only be set as 2 (top-left aligned causal mask) "
-        "or 3 (down-right aligned causal mask)."
+        "Environment variable `NPU_FA2_SPARSE_MODE` can only be set as 3 (down-right aligned causal mask)."
     )
 
 ATTN_MASK_NPU_CACHE = {}
@@ -44,10 +48,6 @@ def get_attn_mask_npu(device):
     return ATTN_MASK_NPU_CACHE[device]
 
 
-def is_npu_fa2_top_left_aligned_causal_mask():
-    return SPARSE_MODE == TOP_LEFT_ALIGNED_CAUSAL_MASK_MODE if is_torch_npu_available() else False
-
-
 def npu_flash_attn_func(
     q,
     k,
@@ -136,8 +136,3 @@ def npu_flash_attn_varlen_func(
         )[0]
 
     return output
-
-
-# This function is not implemented but should never be called because block table is not used on NPU
-def npu_flash_attn_with_kvcache():
-    raise NotImplementedError("npu_flash_attn_with_kvcache is not implemented")