[Question] Handling of custom flex attention block masks
Author: ccdv-ai
Pull request: no

Following the merge of PR #36889 today, I have a question:
How do we deal with custom flex attention functions? Flex attention requires a block mask which is expensive to create.
In practice, the block mask should always be created in the first layer and passed to the next one to avoid any recomputation as seen in the llama implementation:

https://github.com/huggingface/transformers/blob/ad5d40de9c4d4899d5b79243f63e22c72e8b3669/src/transformers/models/llama/modeling_llama.py#L644-L648

@Cyrilvallez 

Comments:
@ccdv-ai Below are two examples demonstrating caching mechanisms to avoid recomputing the block mask in flex attention.
Example 1: Internal Caching Mechanism
In this approach, a cache dictionary is maintained within the model. When the block mask is computed for the first time, it is stored in the cache. For subsequent occurrences of the same input, the result is directly retrieved from the cache without recomputation.

```python
import torch

class FlexAttention:
    def __init__(self):
        # Cache for block masks. The key can be customized; here we use the tensor's data_ptr as an example.
        self.block_mask_cache = {}

    def make_flex_block_causal_mask(self, attention_mask):
        # This function computes the block mask. Replace this example with the actual implementation.
        return attention_mask.clone()

    def get_block_mask(self, attention_mask):
        # Use the memory address of attention_mask as the key for caching.
        key = attention_mask.data_ptr() if isinstance(attention_mask, torch.Tensor) else id(attention_mask)
        if key in self.block_mask_cache:
            # If the block mask is already cached, return it.
            return self.block_mask_cache[key]
        else:
            # Otherwise, compute the block mask and store it in the cache.
            block_mask = self.make_flex_block_causal_mask(attention_mask)
            self.block_mask_cache[key] = block_mask
            return block_mask

# Usage example
attention = FlexAttention()
input_tensor = torch.ones(4, 4)  # Example input

# The first call computes and caches the block mask.
mask1 = attention.get_block_mask(input_tensor)
# Subsequent calls retrieve the block mask directly from the cache.
mask2 = attention.get_block_mask(input_tensor)
print(torch.equal(mask1, mask2))  # Expected output: True
```

Example 2: Lazy Evaluation
This approach delays the computation of the block mask until it is actually needed. Once computed, the result is stored and reused for later calls, preventing redundant calculations.
```python
import torch

class FlexAttentionLazy:
    def __init__(self, attention_mask):
        self._attention_mask = attention_mask
        self._block_mask = None  # Block mask is not computed initially.

    def make_flex_block_causal_mask(self, attention_mask):
        # This function computes the block mask. Replace this with the actual implementation.
        return attention_mask.clone()

    @property
    def block_mask(self):
        # Compute the block mask upon first access.
        if self._block_mask is None:
            self._block_mask = self.make_flex_block_causal_mask(self._attention_mask)
        return self._block_mask

# Usage example
input_tensor = torch.ones(4, 4)  # Example input
attention_lazy = FlexAttentionLazy(input_tensor)

# The first access computes the block mask.
mask_lazy1 = attention_lazy.block_mask
# Later accesses return the already computed block mask.
mask_lazy2 = attention_lazy.block_mask
print(torch.equal(mask_lazy1, mask_lazy2))  # Expected output: True
```
Both methods help avoid repeated computations of the block mask, thereby improving the efficiency of the flex attention implementation. The choice between internal caching and lazy evaluation depends on your model's requirements and specific use cases.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.