206 lines
8.6 KiB
Python
206 lines
8.6 KiB
Python
import torch
|
|
from typing import Callable, Optional
|
|
|
|
|
|
def and_masks(*mask_functions: Callable) -> Callable:
|
|
"""Returns a mask function that is the intersection of provided mask functions"""
|
|
if not all(callable(arg) for arg in mask_functions):
|
|
raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}")
|
|
|
|
def and_mask(batch_idx, head_idx, q_idx, kv_idx):
|
|
result = q_idx.new_ones((), dtype=torch.bool)
|
|
for mask in mask_functions:
|
|
result = result & mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device)
|
|
return result
|
|
|
|
return and_mask
|
|
|
|
|
|
def causal_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
|
|
"""
|
|
This creates a basic lower-diagonal causal mask.
|
|
"""
|
|
return kv_idx <= q_idx
|
|
|
|
|
|
def padding_mask_function(padding_mask: torch.Tensor) -> Callable:
|
|
"""
|
|
This return the mask_function function corresponding to a 2D padding mask.
|
|
"""
|
|
|
|
def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:
|
|
# Note that here the mask should ALWAYS be at least of the max `kv_index` size in the dimension 1. This is because
|
|
# we cannot pad it here in the mask_function as we don't know the final size, and we cannot try/except, as it is not
|
|
# vectorizable on accelerator devices
|
|
return padding_mask[batch_idx, kv_idx]
|
|
|
|
return inner_mask
|
|
|
|
|
|
def prepare_padding_mask(attention_mask: Optional[torch.Tensor], kv_length: int,
|
|
kv_offset: int) -> Optional[torch.Tensor]:
|
|
"""
|
|
From the 2D attention mask, prepare the correct padding mask to use by potentially padding it.
|
|
"""
|
|
local_padding_mask = attention_mask
|
|
if attention_mask is not None:
|
|
# Pad it if necessary
|
|
if (padding_length := kv_length + kv_offset - attention_mask.shape[-1]) > 0:
|
|
local_padding_mask = torch.nn.functional.pad(attention_mask, (0, padding_length))
|
|
return local_padding_mask
|
|
|
|
|
|
def _non_vmap_expansion_sdpa(batch_indices: torch.Tensor, head_indices: torch.Tensor, q_indices: torch.Tensor,
|
|
kv_indices: torch.Tensor):
|
|
"""
|
|
Used to broadcast our mask_functions over the all 4 dimensions (b_idx, h_idx, q_idx, kv_idx) of the inputs.
|
|
Allows the usage of any index-based mask function without relying on vmap.
|
|
|
|
NOTE: This is limited to index based functions only and is not guaranteed to work otherwise.
|
|
|
|
Reference:
|
|
- https://github.com/huggingface/optimum-onnx/blob/c123e8f4fab61b54a8e0e31ce74462bcacca576e/optimum/exporters/onnx/model_patcher.py#L362-L365
|
|
"""
|
|
batch_indices = batch_indices[:, None, None, None]
|
|
head_indices = head_indices[None, :, None, None]
|
|
q_indices = q_indices[None, None, :, None]
|
|
kv_indices = kv_indices[None, None, None, :]
|
|
return batch_indices, head_indices, q_indices, kv_indices
|
|
|
|
|
|
def sdpa_mask(
|
|
batch_size: int,
|
|
cache_position: torch.Tensor,
|
|
kv_length: int,
|
|
kv_offset: int = 0,
|
|
mask_function: Callable = causal_mask_function,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
local_size: Optional[int] = None,
|
|
allow_is_causal_skip: bool = True,
|
|
allow_is_bidirectional_skip: bool = False,
|
|
allow_torch_fix: bool = True,
|
|
use_vmap: bool = False,
|
|
**kwargs,
|
|
) -> Optional[torch.Tensor]:
|
|
"""
|
|
Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that
|
|
the element should take part in the attention computation, and False that it should not.
|
|
This function can only be used with torch>=2.5, as the context manager is otherwise not available.
|
|
|
|
Args:
|
|
batch_size (`int`):
|
|
The batch size of the input sequence.
|
|
cache_position (`torch.Tensor`):
|
|
A tensor of shape (query_length,) indicating the current indices of the input sequence elements.
|
|
kv_length (`int`):
|
|
The size that the key and value states will have during the attention computation.
|
|
kv_offset (`int`, optional):
|
|
An optional offset to indicate at which first position the key and values states will refer to.
|
|
mask_function (`Callable`):
|
|
The mask factory function describing the mask pattern.
|
|
attention_mask (`torch.Tensor`, optional):
|
|
The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)
|
|
local_size (`int`, optional):
|
|
The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True`
|
|
to try to skip mask creation if possible.
|
|
allow_is_causal_skip (`bool`, optional):
|
|
Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in
|
|
`torch.sdpa` instead. Default to `True`.
|
|
allow_is_bidirectional_skip (`bool`, optional):
|
|
Whether to allow to return `None` for the mask under conditions where we do not have to add any bias,
|
|
i.e. full attention without any padding. Default to `False`.
|
|
allow_torch_fix (`bool`, optional):
|
|
Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older
|
|
versions. We need an arg to skip it when using eager. By default `True`.
|
|
use_vmap (`bool`, optional):
|
|
Whether to use `vmap` during the mask construction or not. Allows powerful custom patterns that may not be
|
|
index-based (for the cost of speed performance). By default `False`.
|
|
|
|
|
|
## Creating a simple causal mask:
|
|
|
|
To create the following causal mask:
|
|
|
|
0 ■ ⬚ ⬚ ⬚ ⬚
|
|
1 ■ ■ ⬚ ⬚ ⬚
|
|
2 ■ ■ ■ ⬚ ⬚
|
|
3 ■ ■ ■ ■ ⬚
|
|
4 ■ ■ ■ ■ ■
|
|
|
|
You can do
|
|
|
|
```python
|
|
>>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5)
|
|
>>> tensor([[[[ True, False, False, False, False],
|
|
[ True, True, False, False, False],
|
|
[ True, True, True, False, False],
|
|
[ True, True, True, True, False],
|
|
[ True, True, True, True, True]]]])
|
|
```
|
|
|
|
## Creating a sliding window mask:
|
|
|
|
To create the following sliding window mask (`sliding_window=3`):
|
|
|
|
0 ■ ⬚ ⬚ ⬚ ⬚
|
|
1 ■ ■ ⬚ ⬚ ⬚
|
|
2 ■ ■ ■ ⬚ ⬚
|
|
3 ⬚ ■ ■ ■ ⬚
|
|
4 ⬚ ⬚ ■ ■ ■
|
|
|
|
You can do
|
|
|
|
```python
|
|
>>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=sliding_window_causal_mask_function(3))
|
|
>>> tensor([[[[ True, False, False, False, False],
|
|
[ True, True, False, False, False],
|
|
[ True, True, True, False, False],
|
|
[False, True, True, True, False],
|
|
[False, False, True, True, True]]]])
|
|
```
|
|
|
|
## Creating a chunked attention mask
|
|
|
|
To create the following chunked attention mask (`chunk_size=3`):
|
|
|
|
0 ■ ⬚ ⬚ ⬚ ⬚
|
|
1 ■ ■ ⬚ ⬚ ⬚
|
|
2 ■ ■ ■ ⬚ ⬚
|
|
3 ⬚ ⬚ ⬚ ■ ⬚
|
|
4 ⬚ ⬚ ⬚ ■ ■
|
|
|
|
You can do
|
|
|
|
```python
|
|
>>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=chunked_causal_mask_function(3, torch.zeros(1, dtype=int)))
|
|
>>> tensor([[[[ True, False, False, False, False],
|
|
[ True, True, False, False, False],
|
|
[ True, True, True, False, False],
|
|
[False, False, False, True, False],
|
|
[False, False, False, True, True]]]])
|
|
```
|
|
|
|
"""
|
|
q_length = cache_position.shape[0]
|
|
|
|
# Potentially pad the 2D mask
|
|
padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset)
|
|
|
|
# Potentially add the padding 2D mask
|
|
if padding_mask is not None:
|
|
mask_function = and_masks(mask_function, padding_mask_function(padding_mask))
|
|
|
|
batch_arange = torch.arange(batch_size, device=cache_position.device)
|
|
head_arange = torch.arange(1, device=cache_position.device)
|
|
# Similar to `kv_arange = torch.arange(start=kv_offset, end=kv_offset + kv_length, device=cache_position.device)`
|
|
# but without data-dependent slicing (i.e. torch.compile friendly)
|
|
kv_arange = torch.arange(kv_length, device=cache_position.device) + kv_offset
|
|
|
|
# Actual mask creation
|
|
# Apply mask function element-wise through broadcasting
|
|
attention_mask = mask_function(*_non_vmap_expansion_sdpa(batch_arange, head_arange, cache_position, kv_arange))
|
|
# Expand the mask to match batch size and query length if they weren't used in the mask function
|
|
attention_mask = attention_mask.expand(batch_size, -1, q_length, kv_length)
|
|
|
|
return attention_mask
|