Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2025 HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15import itertools16from typing import Callable, Optional, Union17 18import torch19import torch.nn.functional as F20 21from .cache_utils import Cache22from .configuration_utils import PretrainedConfig23from .utils import is_torch_xpu_available, logging24from .utils.generic import GeneralInterface25from .utils.import_utils import is_torch_flex_attn_available, is_torch_greater_or_equal, is_torchdynamo_compiling26 27 28if is_torch_flex_attn_available():29 from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size30 from torch.nn.attention.flex_attention import BlockMask, create_block_mask31else:32 # Register a fake type to avoid crashing for annotations and `isinstance` checks33 BlockMask = torch.Tensor34 35_is_torch_greater_or_equal_than_2_5 = is_torch_greater_or_equal("2.5", accept_dev=True)36_is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True)37_is_torch_xpu_available = is_torch_xpu_available()38 39if _is_torch_greater_or_equal_than_2_6:40 from torch._dynamo._trace_wrapped_higher_order_op import TransformGetItemToIndex41 42 43logger = logging.get_logger(__name__)44 45 46def and_masks(*mask_functions: Callable) -> Callable:47 """Returns a mask function that is the intersection of provided mask functions"""48 if not all(callable(arg) for arg in mask_functions):49 raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}")50 51 def and_mask(batch_idx, head_idx, q_idx, kv_idx):52 result = q_idx.new_ones((), dtype=torch.bool)53 for mask in mask_functions:54 result = result & mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device)55 return result56 57 return and_mask58 59 60def or_masks(*mask_functions: Callable) -> Callable:61 """Returns a mask function that is the union of provided mask functions"""62 if not all(callable(arg) for arg in mask_functions):63 raise RuntimeError(f"All inputs should be callable mask_functions: {mask_functions}")64 65 def or_mask(batch_idx, head_idx, q_idx, kv_idx):66 result = q_idx.new_zeros((), dtype=torch.bool)67 for mask in mask_functions:68 result = result | mask(batch_idx, head_idx, q_idx, kv_idx).to(result.device)69 return result70 71 return or_mask72 73 74def causal_mask_function(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:75 """76 This creates a basic lower-diagonal causal mask.77 """78 return kv_idx <= q_idx79 80 81def sliding_window_overlay(sliding_window: int) -> Callable:82 """83 This is an overlay depicting a sliding window pattern. Add it on top of a causal mask for a proper sliding84 window mask.85 """86 87 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:88 return kv_idx > q_idx - sliding_window89 90 return inner_mask91 92 93def chunked_overlay(chunk_size: int, left_padding: torch.Tensor) -> Callable:94 """95 This is an overlay depicting a chunked attention pattern. Add it on top of a causal mask for a proper chunked96 attention mask.97 """98 99 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:100 return (kv_idx - left_padding[batch_idx]) // chunk_size == (q_idx - left_padding[batch_idx]) // chunk_size101 102 return inner_mask103 104 105def _legacy_chunked_overlay(chunk_size: int) -> Callable:106 """107 Same as the above function, but do not correctly account for left padding tokens.108 Only kept for compatibility with older torch versions (< 2.6).109 """110 111 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:112 return kv_idx // chunk_size == q_idx // chunk_size113 114 return inner_mask115 116 117def sliding_window_causal_mask_function(sliding_window: int) -> Callable:118 """119 This return the mask_function function to create a sliding window mask.120 """121 return and_masks(sliding_window_overlay(sliding_window), causal_mask_function)122 123 124def chunked_causal_mask_function(chunk_size: int, left_padding: torch.Tensor) -> Callable:125 """126 This return the mask_function function to create a chunked attention mask.127 """128 if not _is_torch_greater_or_equal_than_2_6:129 return and_masks(_legacy_chunked_overlay(chunk_size), causal_mask_function)130 return and_masks(chunked_overlay(chunk_size, left_padding), causal_mask_function)131 132 133def padding_mask_function(padding_mask: torch.Tensor) -> Callable:134 """135 This return the mask_function function corresponding to a 2D padding mask.136 """137 138 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:139 # Note that here the mask should ALWAYS be at least of the max `kv_index` size in the dimension 1. This is because140 # 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 not141 # vectorizable on accelerator devices142 return padding_mask[batch_idx, kv_idx]143 144 return inner_mask145 146 147def packed_sequence_mask_function(packed_sequence_mask: torch.Tensor) -> Callable:148 """149 This return the mask_function function corresponding to a 2D packed sequence mask.150 """151 152 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:153 return packed_sequence_mask[batch_idx, q_idx] == packed_sequence_mask[batch_idx, kv_idx]154 155 return inner_mask156 157 158def add_offsets_to_mask_function(mask_function: Callable, q_offset: int, kv_offset: int) -> Callable:159 """160 This function adds the correct offsets to the `q_idx` and `kv_idx` as the torch API can only accept lengths,161 not start and end indices.162 """163 164 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:165 return mask_function(batch_idx, head_idx, q_idx + q_offset, kv_idx + kv_offset)166 167 return inner_mask168 169 170def _vmap_for_bhqkv(mask_function: Callable, bh_indices: bool = True) -> Callable:171 """172 Used to vmap our mask_functions over the q_idx and kv_idx dimensions of the inputs. Optionally, vmap over173 the batch and head indices as well if `bh_indices=True`.174 Using vmap here allows us to keep the performance of vectorized ops, while having a single set of primitive175 functions between attention interfaces (i.e. between flex and sdpa/eager, FA2 being a bit different).176 177 Args:178 mask_function (`Callable`):179 The mask_function to vmap.180 bh_indices (`bool`, optional):181 Whether to vmap over the batch and head indices as well, or only q and kv indices.182 183 Returns:184 Callable: The vmapped function.185 """186 # We vmap the function 2 times, broadcasting the [q_idx, kv_idx] dimensions187 dimensions = [(None, None, None, 0), (None, None, 0, None)]188 if bh_indices:189 # We extend broadcasting over the [batch_idx, head_idx] dimensions190 dimensions.extend([(None, 0, None, None), (0, None, None, None)])191 192 for dims in dimensions:193 mask_function = torch.vmap(mask_function, in_dims=dims, out_dims=0)194 return mask_function195 196 197def prepare_padding_mask(198 attention_mask: Optional[torch.Tensor], kv_length: int, kv_offset: int, _slice: bool = True199) -> Optional[torch.Tensor]:200 """201 From the 2D attention mask, prepare the correct padding mask to use by potentially padding it, and slicing202 according to the `kv_offset` if `_slice` is `True`.203 """204 local_padding_mask = attention_mask205 if attention_mask is not None:206 # Pad it if necessary207 if (padding_length := kv_length + kv_offset - attention_mask.shape[-1]) > 0:208 local_padding_mask = torch.nn.functional.pad(attention_mask, (0, padding_length))209 # For flex, we should not slice them, only use an offset210 if _slice:211 # Equivalent to: `local_padding_mask = attention_mask[:, kv_offset : kv_offset + kv_length]`,212 # but without data-dependent slicing (i.e. torch.compile friendly)213 mask_indices = torch.arange(kv_length, device=local_padding_mask.device)214 mask_indices += kv_offset215 local_padding_mask = local_padding_mask[:, mask_indices]216 return local_padding_mask217 218 219def _ignore_causal_mask_sdpa(220 padding_mask: Optional[torch.Tensor],221 query_length: int,222 kv_length: int,223 kv_offset: int,224 local_attention_size: Optional[int] = None,225) -> bool:226 """227 Detects whether the causal mask can be ignored in case PyTorch's SDPA is used, rather relying on SDPA's `is_causal` argument.228 229 In case no token is masked in the 2D `padding_mask` argument, if `query_length == 1` or230 `key_value_length == query_length`, we rather rely on SDPA `is_causal` argument to use causal/non-causal masks,231 allowing to dispatch to the flash attention kernel (that can otherwise not be used if a custom `attn_mask` is232 passed).233 """234 is_tracing = torch.jit.is_tracing() or isinstance(padding_mask, torch.fx.Proxy) or is_torchdynamo_compiling()235 if padding_mask is not None and padding_mask.shape[-1] > kv_length:236 mask_indices = torch.arange(kv_length, device=padding_mask.device)237 mask_indices += kv_offset238 padding_mask = padding_mask[:, mask_indices]239 240 # When using `torch.export` or `torch.onnx.dynamo_export`, we must pass an example input, and `is_causal` behavior is241 # hard-coded to the forward. If a user exports a model with query_length > 1, the exported model will hard-code `is_causal=True`242 # which is in general wrong (see https://github.com/pytorch/pytorch/issues/108108). Thus, we only set243 # `ignore_causal_mask = True` if we are not tracing244 if (245 not is_tracing246 # only cases when lower and upper diags are the same, see https://github.com/pytorch/pytorch/issues/108108247 and (query_length == 1 or (kv_length == query_length or _is_torch_xpu_available))248 # in this case we need to add special patterns to the mask so cannot be skipped otherwise249 and (local_attention_size is None or kv_length < local_attention_size)250 # In this case, we need to add padding to the mask, so cannot be skipped otherwise251 and (252 padding_mask is None253 or (254 padding_mask.all()255 if not _is_torch_xpu_available or query_length == 1256 else padding_mask[:, :query_length].all()257 )258 )259 ):260 return True261 262 return False263 264 265def sdpa_mask_recent_torch(266 batch_size: int,267 cache_position: torch.Tensor,268 kv_length: int,269 kv_offset: int = 0,270 mask_function: Callable = causal_mask_function,271 attention_mask: Optional[torch.Tensor] = None,272 local_size: Optional[int] = None,273 allow_is_causal_skip: bool = True,274 **kwargs,275) -> Optional[torch.Tensor]:276 """277 Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that278 the element should take part in the attention computation, and False that it should not.279 This function can only be used with torch>=2.5, as the context manager is otherwise not available.280 281 Args:282 batch_size (`int`):283 The batch size of the input sequence.284 cache_position (`torch.Tensor`):285 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.286 kv_length (`int`):287 The size that the key and value states will have during the attention computation.288 kv_offset (`int`, optional):289 An optional offset to indicate at which first position the key and values states will refer to.290 mask_function (`Callable`):291 The mask factory function describing the mask pattern.292 attention_mask (`torch.Tensor`, optional):293 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)294 local_size (`int`, optional):295 The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True`296 to try to skip mask creation if possible.297 allow_is_causal_skip (`bool`, optional):298 Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in299 `torch.sdpa` instead. Default to `True`.300 allow_torch_fix (`bool`, optional):301 Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older302 versions. We need an arg to skip it when using eager. By default `True`.303 304 305 ## Creating a simple causal mask:306 307 To create the following causal mask:308 309 0 ■ ⬚ ⬚ ⬚ ⬚310 1 ■ ■ ⬚ ⬚ ⬚311 2 ■ ■ ■ ⬚ ⬚312 3 ■ ■ ■ ■ ⬚313 4 ■ ■ ■ ■ ■314 315 You can do316 317 ```python318 >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5)319 >>> tensor([[[[ True, False, False, False, False],320 [ True, True, False, False, False],321 [ True, True, True, False, False],322 [ True, True, True, True, False],323 [ True, True, True, True, True]]]])324 ```325 326 ## Creating a sliding window mask:327 328 To create the following sliding window mask (`sliding_window=3`):329 330 0 ■ ⬚ ⬚ ⬚ ⬚331 1 ■ ■ ⬚ ⬚ ⬚332 2 ■ ■ ■ ⬚ ⬚333 3 ⬚ ■ ■ ■ ⬚334 4 ⬚ ⬚ ■ ■ ■335 336 You can do337 338 ```python339 >>> sdpa_mask(batch_size=1, cache_position=torch.arange(5), kv_length=5, mask_function=sliding_window_causal_mask_function(3))340 >>> tensor([[[[ True, False, False, False, False],341 [ True, True, False, False, False],342 [ True, True, True, False, False],343 [False, True, True, True, False],344 [False, False, True, True, True]]]])345 ```346 347 ## Creating a chunked attention mask348 349 To create the following chunked attention mask (`chunk_size=3`):350 351 0 ■ ⬚ ⬚ ⬚ ⬚352 1 ■ ■ ⬚ ⬚ ⬚353 2 ■ ■ ■ ⬚ ⬚354 3 ⬚ ⬚ ⬚ ■ ⬚355 4 ⬚ ⬚ ⬚ ■ ■356 357 You can do358 359 ```python360 >>> 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)))361 >>> tensor([[[[ True, False, False, False, False],362 [ True, True, False, False, False],363 [ True, True, True, False, False],364 [False, False, False, True, False],365 [False, False, False, True, True]]]])366 ```367 368 """369 q_length = cache_position.shape[0]370 # Potentially pad the 2D mask, and slice it correctly371 padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False)372 373 # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument374 if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size):375 return None376 377 # Similar to `kv_arange = torch.arange(start=kv_offset, end=kv_offset + kv_length, device=cache_position.device)`378 # but without data-dependent slicing (i.e. torch.compile friendly)379 kv_arange = torch.arange(kv_length, device=cache_position.device)380 kv_arange += kv_offset381 382 # Potentially add the padding 2D mask383 if padding_mask is not None:384 mask_function = and_masks(mask_function, padding_mask_function(padding_mask))385 386 batch_arange = torch.arange(batch_size, device=cache_position.device)387 head_arange = torch.arange(1, device=cache_position.device)388 # This creates the 4D mask easily. Note that we need this context manager as vmap cannot handle slicing a tensor from389 # scalar tensor (it internally calls `.item()` which vmap does not allow, but this context works around it390 # We don't need to add an offset to the mask_function either, as we vmap directly the correct indices for k and kv indices391 with TransformGetItemToIndex():392 causal_mask = _vmap_for_bhqkv(mask_function)(batch_arange, head_arange, cache_position, kv_arange)393 394 return causal_mask395 396 397def sdpa_mask_older_torch(398 batch_size: int,399 cache_position: torch.Tensor,400 kv_length: int,401 kv_offset: int = 0,402 mask_function: Callable = causal_mask_function,403 attention_mask: Optional[torch.Tensor] = None,404 local_size: Optional[int] = None,405 allow_is_causal_skip: bool = True,406 allow_torch_fix: bool = True,407 **kwargs,408) -> Optional[torch.Tensor]:409 """410 NOTE: This function is only used when torch version is torch<2.5 - see `sdpa_mask_recent_torch` otherwise.411 412 Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that413 the element should take part in the attention computation, and False that it should not.414 If `allow_torch_fix=True` (the default), rows corresponding to query tokens that do not attend415 to any other tokens (due to padding) will be fully attended to instead, in order to avoid `nan` propagation (this does416 not change the final result).417 418 Args:419 batch_size (`int`):420 The batch size of the input sequence.421 cache_position (`torch.Tensor`):422 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.423 kv_length (`int`):424 The size that the key and value states will have during the attention computation.425 kv_offset (`int`, optional):426 An optional offset to indicate at which first position the key and values states will refer to.427 mask_function (`Callable`):428 The mask factory function describing the mask pattern.429 attention_mask (`torch.Tensor`, optional):430 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)431 local_size (`int`, optional):432 The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True`433 to try to skip mask creation if possible.434 allow_is_causal_skip (`bool`, optional):435 Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in436 `torch.sdpa` instead. Default to `True`.437 allow_torch_fix (`bool`, optional):438 Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older439 versions. We need an arg to skip it when using eager. By default `True`.440 """441 q_length = cache_position.shape[0]442 # Potentially pad the 2D mask, and slice it correctly443 padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset)444 445 # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument446 if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size):447 return None448 449 # Similar to `kv_arange = torch.arange(start=kv_offset, end=kv_offset + kv_length, device=cache_position.device)`450 # but without data-dependent slicing (i.e. torch.compile friendly)451 kv_arange = torch.arange(kv_length, device=cache_position.device)452 kv_arange += kv_offset453 454 # This creates the 4D mask easily. Note that we do not include vmap over the batch_idx dimension as well,455 # as vmap cannot handle slicing a tensor from scalar tensor (it internally calls `.item()` which vmap does not allow456 # However, in more recent version of Pytorch, a trick was introduced to handle it - which is the reason we have457 # `sdpa_mask_recent_torch`, as it allows more general `mask_function`458 causal_mask = _vmap_for_bhqkv(mask_function, bh_indices=False)(None, None, cache_position, kv_arange)459 causal_mask = causal_mask[None, None, :, :].expand(batch_size, -1, -1, -1)460 if padding_mask is not None:461 causal_mask = causal_mask * padding_mask[:, None, None, :]462 463 # Due to a bug in versions of torch<2.5, we need to update the mask in case a query is not attending to any464 # tokens (due to padding). See details in https://github.com/pytorch/pytorch/issues/110213465 if not _is_torch_greater_or_equal_than_2_5 and allow_torch_fix:466 causal_mask |= torch.all(~causal_mask, dim=-1, keepdim=True)467 return causal_mask468 469 470# We use the version with newer torch whenever possible, as it is more general and can handle arbitrary mask functions471# (especially mask_function indexing a tensor, such as the padding mask function)472sdpa_mask = sdpa_mask_recent_torch if _is_torch_greater_or_equal_than_2_6 else sdpa_mask_older_torch473 474 475def eager_mask(476 batch_size: int,477 cache_position: torch.Tensor,478 kv_length: int,479 kv_offset: int = 0,480 mask_function: Callable = causal_mask_function,481 attention_mask: Optional[torch.Tensor] = None,482 dtype: torch.dtype = torch.float32,483 **kwargs,484) -> torch.Tensor:485 """486 Create a 4D float mask of shape `(batch_size, 1, query_length, kv_length)` where a value of 0 indicates that487 the element should take part in the attention computation, and -inf (minimum value for the given `dtype`) that488 it should not.489 490 Args:491 batch_size (`int`):492 The batch size of the input sequence.493 cache_position (`torch.Tensor`):494 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.495 kv_length (`int`):496 The size that the key and value states will have during the attention computation.497 kv_offset (`int`, optional):498 An optional offset to indicate at which first position the key and values states will refer to.499 mask_function (`Callable`):500 The mask factory function describing the mask pattern.501 attention_mask (`torch.Tensor`, optional):502 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)503 dtype (`torch.dtype`, optional):504 The dtype to use for the mask. By default, `torch.float32`.505 """506 # The masks for eager attention are simply boolean mask from sdpa, casted to 0 and -inf507 _ = kwargs.pop("allow_is_causal_skip", None)508 mask = sdpa_mask(509 batch_size=batch_size,510 cache_position=cache_position,511 kv_length=kv_length,512 kv_offset=kv_offset,513 mask_function=mask_function,514 attention_mask=attention_mask,515 allow_is_causal_skip=False,516 allow_torch_fix=False,517 **kwargs,518 )519 min_dtype = torch.finfo(dtype).min520 # we need 0s where the tokens should be taken into account, and -inf otherwise (mask is already of boolean type)521 mask = torch.where(mask, torch.tensor(0.0, device=mask.device, dtype=dtype), min_dtype)522 return mask523 524 525def flash_attention_mask(526 batch_size: int,527 cache_position: torch.Tensor,528 kv_length: int,529 kv_offset: int = 0,530 mask_function: Callable = causal_mask_function,531 attention_mask: Optional[torch.Tensor] = None,532 **kwargs,533):534 """535 Create the attention mask necessary to use FA2. Since FA2 is un-padded by definition, here we simply return536 `None` if the mask is fully causal, or we return the 2D mask which will then be used to extract the seq_lens.537 We just slice it in case of sliding window.538 539 Args:540 batch_size (`int`):541 The batch size of the input sequence.542 cache_position (`torch.Tensor`):543 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.544 kv_length (`int`):545 The size that the key and value states will have during the attention computation.546 kv_offset (`int`, optional):547 An optional offset to indicate at which first position the key and values states will refer to.548 mask_function (`Callable`):549 The mask factory function describing the mask pattern.550 attention_mask (`torch.Tensor`, optional):551 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)552 """553 if attention_mask is not None:554 # Here we need to slice from the right if using sliding or chunked (for full attention, this is equivalent to doing nothing)555 attention_mask = attention_mask[:, -kv_length:]556 # We only return an actual mask if there is at least 1 padding token, otherwise we return `None` and use `is_causal` in FA2557 # (note that the attention_mask is a boolean dtype here)558 if attention_mask.all():559 attention_mask = None560 561 return attention_mask562 563 564def flex_attention_mask(565 batch_size: int,566 cache_position: torch.Tensor,567 kv_length: int,568 kv_offset: int = 0,569 mask_function: Callable = causal_mask_function,570 attention_mask: Optional[torch.Tensor] = None,571 **kwargs,572) -> BlockMask:573 """574 Create a 4D block mask which is a compressed representation of the full 4D block causal mask. BlockMask is essential575 for performant computation of flex attention. See: https://pytorch.org/blog/flexattention/576 577 Args:578 batch_size (`int`):579 The batch size of the input sequence.580 cache_position (`torch.Tensor`):581 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.582 kv_length (`int`):583 The size that the key and value states will have during the attention computation.584 kv_offset (`int`, optional):585 An optional offset to indicate at which first position the key and values states will refer to.586 mask_function (`Callable`):587 The mask factory function describing the mask pattern.588 attention_mask (`torch.Tensor`, optional):589 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)590 """591 q_length, q_offset = cache_position.shape[0], cache_position[0]592 593 # Potentially add the padding 2D mask594 if attention_mask is not None:595 # Older torch (2.5.x) cannot handle sequences not in multiples of 128 (default block size)596 # Hence we pad to multiples of this as a minimum to ensure this597 pad_len = ((attention_mask.shape[1] // flex_default_block_size) + 1) * flex_default_block_size598 pad_len = pad_len - attention_mask.shape[1]599 if not _is_torch_greater_or_equal_than_2_6 and pad_len > 0:600 attention_mask = torch.nn.functional.pad(attention_mask, value=0, pad=(0, pad_len))601 602 padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False)603 mask_function = and_masks(mask_function, padding_mask_function(padding_mask))604 605 # Add the offsets on top (because flex interface only allows length, not start and end indices)606 mask_function = add_offsets_to_mask_function(mask_function, q_offset, kv_offset)607 608 # Finally create the block mask609 block_mask = create_block_mask(610 mask_mod=mask_function,611 B=batch_size,612 H=None,613 Q_LEN=q_length,614 KV_LEN=kv_length,615 device=cache_position.device,616 _compile=_is_torch_greater_or_equal_than_2_6,617 )618 return block_mask619 620 621class AttentionMaskInterface(GeneralInterface):622 # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if623 # a new instance is created (in order to locally override a given function)624 _global_mapping = {625 "sdpa": sdpa_mask,626 "eager": eager_mask,627 "flash_attention_2": flash_attention_mask,628 "flash_attention_3": flash_attention_mask,629 "flex_attention": flex_attention_mask,630 }631 632 633# Global AttentionMaskInterface shared by all models which do not need to overwrite any of the existing ones634ALL_MASK_ATTENTION_FUNCTIONS: AttentionMaskInterface = AttentionMaskInterface()635 636 637def find_packed_sequence_indices(position_ids: torch.Tensor) -> torch.Tensor:638 """639 Find the indices of the sequence to which each new query token in the sequence belongs when using packed640 tensor format (i.e. several sequences packed in the same batch dimension).641 642 Args:643 position_ids (`torch.Tensor`)644 A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.645 646 Returns:647 A 2D tensor where each similar integer indicates that the tokens belong to the same sequence. For example, if we648 pack 3 sequences of 2, 3 and 1 tokens respectively along a single batch dim, this will return [[0, 0, 1, 1, 1, 2]].649 """650 # What separate different sequences is when 2 consecutive positions_ids are separated by more than 1. So651 # taking the diff (by prepending the first value - 1 to keep correct indexing) and applying cumsum to the result652 # gives exactly the sequence indices653 # Note that we assume that a single sequence cannot span several batch dimensions, i.e. 1 single sequence654 # cannot be part of the end of the first batch dim and the start of the 2nd one for example655 first_dummy_value = position_ids[:, :1] - 1 # We just need the diff on this first value to be 1656 position_diff = torch.diff(position_ids, prepend=first_dummy_value, dim=-1)657 packed_sequence_mask = (position_diff != 1).cumsum(-1)658 659 # Here it would be nice to return None if we did not detect packed sequence format, i.e. if `packed_sequence_mask[:, -1] == 0`660 # but it causes issues with export661 return packed_sequence_mask662 663 664def _preprocess_mask_arguments(665 config: PretrainedConfig,666 input_embeds: torch.Tensor,667 attention_mask: Optional[Union[torch.Tensor, BlockMask]],668 cache_position: torch.Tensor,669 past_key_values: Optional[Cache],670 position_ids: Optional[torch.Tensor],671 layer_idx: Optional[int],672) -> tuple[bool, Optional[Union[torch.Tensor, BlockMask]], int, int]:673 """674 Perform some common pre-processing of the mask arguments we get from the modeling code. Mostly determine the675 key-value length and offsets, and if we should early exit or not.676 677 Args:678 config (`PretrainedConfig`):679 The model config.680 input_embeds (`torch.Tensor`):681 The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the682 batch size, query length and dtype.683 attention_mask (`torch.Tensor`, optional):684 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).685 It can also be an already prepared 4D mask, in which case it is returned as-is.686 cache_position (`torch.Tensor`):687 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.688 past_key_values (`Cache`, optional):689 The past key values, if we use a cache.690 position_ids (`torch.Tensor`, optional)691 A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.692 layer_idx (`int`, optional):693 If `past_key_values` is not None, this is the layer index of the cache from which to get the key-value694 length and offset. Indeed, for hybrid caches, different layers may return different lengths.695 696 Returns:697 early_exit (`bool`):698 Whether we should early exit mask creation, and return the mask as-is.699 attention_mask (`torch.Tensor` or `BlockMask` or `None`):700 The attention mask to either return immediately, or to use in downstream mask creation.701 packed_sequence_mask (`torch.Tensor`, optional):702 In case we detected packed sequence format, this is a tensor where each similar integer indicates that703 the tokens belong to the same sequence.704 kv_length (`int`):705 The size that the key and value states will have during the attention computation.706 kv_offset (`int`):707 An offset to indicate at which first position the key and values states will refer to.708 """709 # If the mask is already 4D, simply return as-is (it was already prepared, or it is custom)710 if isinstance(attention_mask, (torch.Tensor, BlockMask)) and len(attention_mask.shape) == 4:711 return True, attention_mask, None, None, None712 713 # For TGI/vLLM backends, or other custom attention without equivalent mask creation: we don't need a mask!714 # Note: it's not ideal to check the `_global_mapping` attribute instead of the object itself, however otherwise715 # full graph dynamo tracing (i.e. torch.export or compile with `fullgraph=True`) will fail on Python<3.11716 # with `torch._dynamo.exc.Unsupported: 'inline in skipfiles:Mapping.__contains__ | __contains__, skipped717 # according trace_rules.lookup SKIP_DIRS'` -- can be removed when we require Python>=3.11718 if config._attn_implementation not in ALL_MASK_ATTENTION_FUNCTIONS._global_mapping:719 return True, None, None, None, None720 721 # Move the mask to correct device, and potentially switch dtype for efficiency722 if attention_mask is not None and attention_mask.ndim == 2:723 attention_mask = attention_mask.to(device=cache_position.device, dtype=torch.bool)724 725 # If using a cache, it can give all information about mask sizes based on seen tokens726 if past_key_values is not None:727 kv_length, kv_offset = past_key_values.get_mask_sizes(cache_position, layer_idx)728 # Otherwise, the sizes are simply the input sizes729 else:730 kv_length, kv_offset = input_embeds.shape[1], 0731 732 # We check the position_ids for potential packed sequence format (only if the 2D attention mask is explicitly None,733 # and we don't have past_key_values, i.e. generally a training setup)734 packed_sequence_mask = None735 if position_ids is not None and attention_mask is None and past_key_values is None:736 batch_size = input_embeds.shape[0]737 # The position ids are sometimes just unsqueezed, without being expanded738 if batch_size != position_ids.shape[0]:739 position_ids = position_ids.expand(batch_size, -1)740 packed_sequence_mask = find_packed_sequence_indices(position_ids)741 742 return False, attention_mask, packed_sequence_mask, kv_length, kv_offset743 744 745def create_causal_mask(746 config: PretrainedConfig,747 input_embeds: torch.Tensor,748 attention_mask: Optional[torch.Tensor],749 cache_position: torch.Tensor,750 past_key_values: Optional[Cache],751 position_ids: Optional[torch.Tensor] = None,752 or_mask_function: Optional[Callable] = None,753 and_mask_function: Optional[Callable] = None,754) -> Optional[Union[torch.Tensor, BlockMask]]:755 """756 Create a standard causal mask based on the attention implementation used (stored in the config). If `past_key_values`757 has an hybrid cache structure, this function will return the mask corresponding to one of the "full_attention" layers (to align758 to what is needed in the `modeling_xxx.py` files).759 760 Args:761 config (`PretrainedConfig`):762 The model config.763 input_embeds (`torch.Tensor`):764 The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the765 batch size, query length and dtype.766 attention_mask (`torch.Tensor`, optional):767 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).768 It can also be an already prepared 4D mask, in which case it is returned as-is.769 cache_position (`torch.Tensor`):770 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.771 past_key_values (`Cache`, optional):772 The past key values, if we use a cache.773 position_ids (`torch.Tensor`, optional)774 A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.775 or_mask_function (`Callable`, optional):776 An optional mask function to combine with the causal mask function (by doing the union of both). This is777 useful to easily overlay another mask on top of the causal one, for example for image tokens handling.778 and_mask_function (`Callable`, optional):779 An optional mask function to combine with the causal mask function (by doing the intersection of both). This is780 useful to easily overlay another mask on top of the causal one, for example for image tokens handling.781 """782 # If we have an hybrid cache structure, here we want to create the mask for the full layers783 if hasattr(past_key_values, "is_sliding") and False in past_key_values.is_sliding:784 layer_idx = past_key_values.is_sliding.index(False)785 else:786 layer_idx = 0787 788 early_exit, attention_mask, packed_sequence_mask, kv_length, kv_offset = _preprocess_mask_arguments(789 config, input_embeds, attention_mask, cache_position, past_key_values, position_ids, layer_idx790 )791 if early_exit:792 return attention_mask793 794 batch_size, dtype = input_embeds.shape[0], input_embeds.dtype795 mask_factory_function = causal_mask_function796 mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]797 798 # Do not allow skip if we are compiling (this is to match BC)799 # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it800 if _is_torch_xpu_available:801 allow_is_causal_skip = True802 else:803 allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False)804 805 # Allow slight deviations from causal mask806 # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,807 # padding mask, etc) as the resulting mask may otherwise not be correct!808 if or_mask_function is not None:809 if not _is_torch_greater_or_equal_than_2_6:810 raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")811 mask_factory_function = or_masks(mask_factory_function, or_mask_function)812 allow_is_causal_skip = False813 if and_mask_function is not None:814 if not _is_torch_greater_or_equal_than_2_6:815 raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")816 mask_factory_function = and_masks(mask_factory_function, and_mask_function)817 allow_is_causal_skip = False818 819 # If we detected packing format820 if packed_sequence_mask is not None and _is_torch_greater_or_equal_than_2_6:821 mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask))822 allow_is_causal_skip = False823 824 # We now create the mask825 causal_mask = mask_interface(826 batch_size=batch_size,827 cache_position=cache_position,828 kv_length=kv_length,829 kv_offset=kv_offset,830 mask_function=mask_factory_function,831 attention_mask=attention_mask,832 allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa833 dtype=dtype, # Additional kwarg for eager834 config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface835 )836 return causal_mask837 838 839def create_sliding_window_causal_mask(840 config: PretrainedConfig,841 input_embeds: torch.Tensor,842 attention_mask: Optional[torch.Tensor],843 cache_position: torch.Tensor,844 past_key_values: Optional[Cache],845 position_ids: Optional[torch.Tensor] = None,846 or_mask_function: Optional[Callable] = None,847 and_mask_function: Optional[Callable] = None,848) -> Optional[Union[torch.Tensor, BlockMask]]:849 """850 Create a sliding window causal mask based on the attention implementation used (stored in the config). This type851 of attention pattern was mostly democratized by Mistral. If `past_key_values` has an hybrid cache structure, this852 function will return the mask corresponding to one of the "sliding_attention" layers (to align to what is needed in the853 `modeling_xxx.py` files).854 855 Args:856 config (`PretrainedConfig`):857 The model config.858 input_embeds (`torch.Tensor`):859 The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the860 batch size, query length and dtype.861 attention_mask (`torch.Tensor`, optional):862 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).863 It can also be an already prepared 4D mask, in which case it is returned as-is.864 cache_position (`torch.Tensor`):865 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.866 past_key_values (`Cache`, optional):867 The past key values, if we use a cache.868 position_ids (`torch.Tensor`, optional)869 A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.870 or_mask_function (`Callable`, optional):871 An optional mask function to combine with the sliding causal mask function (by doing the union of both). This is872 useful to easily overlay another mask on top of the sliding causal one, for example for image tokens handling.873 and_mask_function (`Callable`, optional):874 An optional mask function to combine with the sliding causal mask function (by doing the intersection of both). This is875 useful to easily overlay another mask on top of the sliding causal one, for example for image tokens handling.876 """877 # If we have an hybrid cache structure, here we want to create the mask for the sliding layers878 if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding:879 layer_idx = past_key_values.is_sliding.index(True)880 else:881 layer_idx = 0882 883 early_exit, attention_mask, packed_sequence_mask, kv_length, kv_offset = _preprocess_mask_arguments(884 config, input_embeds, attention_mask, cache_position, past_key_values, position_ids, layer_idx885 )886 if early_exit:887 return attention_mask888 889 sliding_window = getattr(config, "sliding_window", None)890 if sliding_window is None:891 raise ValueError("Could not find a `sliding_window` argument in the config, or it is not set")892 893 batch_size, dtype = input_embeds.shape[0], input_embeds.dtype894 mask_factory_function = sliding_window_causal_mask_function(sliding_window)895 mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]896 897 # Do not allow skip if we are compiling (this is to match BC)898 # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it899 allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False)900 901 # Allow slight deviations from causal mask902 # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,903 # padding mask, etc) as the resulting mask may otherwise not be correct!904 if or_mask_function is not None:905 if not _is_torch_greater_or_equal_than_2_6:906 raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")907 mask_factory_function = or_masks(mask_factory_function, or_mask_function)908 allow_is_causal_skip = False909 if and_mask_function is not None:910 if not _is_torch_greater_or_equal_than_2_6:911 raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")912 mask_factory_function = and_masks(mask_factory_function, and_mask_function)913 allow_is_causal_skip = False914 915 # If we detected packing format916 if packed_sequence_mask is not None and _is_torch_greater_or_equal_than_2_6:917 mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask))918 allow_is_causal_skip = False919 920 # We now create the mask921 causal_mask = mask_interface(922 batch_size=batch_size,923 cache_position=cache_position,924 kv_length=kv_length,925 kv_offset=kv_offset,926 mask_function=mask_factory_function,927 attention_mask=attention_mask,928 allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa929 local_size=sliding_window, # Additional kwarg for sdpa930 dtype=dtype, # Additional kwarg for eager931 config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface932 )933 return causal_mask934 935 936def create_chunked_causal_mask(937 config: PretrainedConfig,938 input_embeds: torch.Tensor,939 attention_mask: Optional[torch.Tensor],940 cache_position: torch.Tensor,941 past_key_values: Optional[Cache],942 position_ids: Optional[torch.Tensor] = None,943 or_mask_function: Optional[Callable] = None,944 and_mask_function: Optional[Callable] = None,945) -> Optional[Union[torch.Tensor, BlockMask]]:946 """947 Create a chunked attention causal mask based on the attention implementation used (stored in the config). This type948 of attention pattern was mostly democratized by Llama4. If `past_key_values` has an hybrid cache structure, this949 function will return the mask corresponding to one of the "chunked_attention" layers (to align to what is needed in the950 `modeling_xxx.py` files).951 952 Args:953 config (`PretrainedConfig`):954 The model config.955 input_embeds (`torch.Tensor`):956 The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the957 batch size, query length and dtype.958 attention_mask (`torch.Tensor`, optional):959 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).960 It can also be an already prepared 4D mask, in which case it is returned as-is.961 cache_position (`torch.Tensor`):962 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.963 past_key_values (`Cache`, optional):964 The past key values, if we use a cache.965 position_ids (`torch.Tensor`, optional)966 A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.967 or_mask_function (`Callable`, optional):968 An optional mask function to combine with the chunked causal mask function (by doing the union of both). This is969 useful to easily overlay another mask on top of the chunked causal one, for example for image tokens handling.970 and_mask_function (`Callable`, optional):971 An optional mask function to combine with the chunked causal mask function (by doing the intersection of both). This is972 useful to easily overlay another mask on top of the chunked causal one, for example for image tokens handling.973 """974 # If we have an hybrid cache structure, here we want to create the mask for the sliding layers975 if hasattr(past_key_values, "is_sliding") and True in past_key_values.is_sliding:976 layer_idx = past_key_values.is_sliding.index(True)977 else:978 layer_idx = 0979 980 early_exit, attention_mask, packed_sequence_mask, kv_length, kv_offset = _preprocess_mask_arguments(981 config, input_embeds, attention_mask, cache_position, past_key_values, position_ids, layer_idx982 )983 if early_exit:984 return attention_mask985 986 chunk_size = getattr(config, "attention_chunk_size", None)987 if chunk_size is None:988 raise ValueError("Could not find an `attention_chunk_size` argument in the config, or it is not set")989 990 # Raise if using chunked attention on context too large with FA2991 if config._attn_implementation == "flash_attention_2" and kv_length + kv_offset > chunk_size:992 raise ValueError(993 "Flash attention 2 cannot handle chunked attention, and the key-value length is larger than the chunk size so the "994 "chunked pattern cannot be respected. You should use another `attn_implementation` when instantiating the model"995 )996 997 batch_size, dtype = input_embeds.shape[0], input_embeds.dtype998 # For chunked attention and batched inputs, we need to take the number of left padding tokens into account999 # to start the chunk from the actual start of the sequence for the padded sequence1000 if attention_mask is not None:1001 # Only count the left padding tokens, not all of them1002 left_padding_tokens = (attention_mask.cumsum(dim=-1) == torch.zeros_like(attention_mask)).sum(dim=-1)1003 else:1004 left_padding_tokens = torch.zeros(batch_size, device=cache_position.device, dtype=int)1005 # Raise a warning for older versions if the problematic left-padding situation arises1006 if (1007 not _is_torch_greater_or_equal_than_2_61008 and kv_length + kv_offset > chunk_size1009 and (left_padding_tokens > 0).any()1010 ):1011 logger.warning_once(1012 "Due to limitations of your current torch version, we cannot correctly account for the left-padding "1013 "when computing the chunked attention pattern. This will lead to a wrong attention mask for the padded "1014 "sequences. Behavior will be undefined. Please upgrade to `torch>=2.6` to solve this issue."1015 )1016 mask_factory_function = chunked_causal_mask_function(chunk_size, left_padding_tokens)1017 mask_interface = ALL_MASK_ATTENTION_FUNCTIONS[config._attn_implementation]1018 1019 # Do not allow skip if we are compiling (this is to match BC)1020 # TODO: cyril -> probably revisit and remove this, but a lot of tests rely on it1021 allow_is_causal_skip = not getattr(past_key_values, "is_compileable", False)1022 1023 # Allow slight deviations from causal mask1024 # Note that it is very important to apply this before any other deviations of the mask (such as packed sequence mask,1025 # padding mask, etc) as the resulting mask may otherwise not be correct!1026 if or_mask_function is not None:1027 if not _is_torch_greater_or_equal_than_2_6:1028 raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")1029 mask_factory_function = or_masks(mask_factory_function, or_mask_function)1030 allow_is_causal_skip = False1031 if and_mask_function is not None:1032 if not _is_torch_greater_or_equal_than_2_6:1033 raise ValueError("Using `or_mask_function` or `and_mask_function` arguments require torch>=2.6")1034 mask_factory_function = and_masks(mask_factory_function, and_mask_function)1035 allow_is_causal_skip = False1036 1037 # If we detected packing format1038 if packed_sequence_mask is not None and _is_torch_greater_or_equal_than_2_6:1039 mask_factory_function = and_masks(mask_factory_function, packed_sequence_mask_function(packed_sequence_mask))1040 allow_is_causal_skip = False1041 1042 # We now create the mask1043 causal_mask = mask_interface(1044 batch_size=batch_size,1045 cache_position=cache_position,1046 kv_length=kv_length,1047 kv_offset=kv_offset,1048 mask_function=mask_factory_function,1049 attention_mask=attention_mask,1050 allow_is_causal_skip=allow_is_causal_skip, # additional kwarg for sdpa1051 local_size=chunk_size, # Additional kwarg for sdpa1052 dtype=dtype, # Additional kwarg for eager1053 config=config, # Pass the config as well, in case someone wants to easily have their own mask_interface1054 )1055 return causal_mask1056 1057 1058LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING = {1059 "full_attention": create_causal_mask,1060 "sliding_attention": create_sliding_window_causal_mask,1061 "chunked_attention": create_chunked_causal_mask,1062}1063 1064 1065def create_masks_for_generate(1066 config: PretrainedConfig,1067 input_embeds: torch.Tensor,1068 attention_mask: Optional[torch.Tensor],1069 cache_position: torch.Tensor,1070 past_key_values: Optional[Cache],1071 position_ids: Optional[torch.Tensor] = None,1072 or_mask_function: Optional[Callable] = None,1073 and_mask_function: Optional[Callable] = None,1074 **kwargs,1075):1076 """1077 This function mimics how we create the masks in the `modeling_xxx.py` files, and is used in `generate` in order1078 to easily create the masks in advance, when we compile the forwards with Static caches.1079 1080 Args:1081 config (`PretrainedConfig`):1082 The model config.1083 input_embeds (`torch.Tensor`):1084 The input embeddings of shape (batch_size, query_length, hidden_dim). This is used only to infer the1085 batch size, query length and dtype.1086 attention_mask (`torch.Tensor`, optional):1087 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length).1088 It can also be an already prepared 4D mask, in which case it is returned as-is.1089 cache_position (`torch.Tensor`):1090 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.1091 past_key_values (`Cache`, optional):1092 The past key values, if we use a cache.1093 position_ids (`torch.Tensor`, optional)1094 A 2D tensor of shape (batch_size, query_length) indicating the positions of each token in the sequences.1095 or_mask_function (`Callable`, optional):1096 An optional mask function to combine with the other mask function (by doing the union of both). This is1097 useful to easily overlay another mask on top of the causal one, for example for image tokens handling.1098 and_mask_function (`Callable`, optional):1099 An optional mask function to combine with the other mask function (by doing the intersection of both). This is1100 useful to easily overlay another mask on top of the causal one, for example for image tokens handling.1101 """1102 # The attribute reside in the text config for composite models1103 effective_config = config.get_text_config()1104 # Prepare the mask args1105 mask_kwargs = {1106 "config": effective_config,1107 "input_embeds": input_embeds,1108 "attention_mask": attention_mask,1109 "cache_position": cache_position,1110 "past_key_values": past_key_values,1111 "position_ids": position_ids,1112 "or_mask_function": or_mask_function,1113 "and_mask_function": and_mask_function,1114 }1115 1116 # If the attribute exist, we need several masks1117 if hasattr(effective_config, "layer_types"):1118 causal_masks = {}1119 for layer_pattern in set(effective_config.layer_types):1120 causal_masks[layer_pattern] = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[layer_pattern](**mask_kwargs)1121 return causal_masks1122 # In this case, all layers are sliding1123 elif getattr(effective_config, "sliding_window", None) is not None:1124 return create_sliding_window_causal_mask(**mask_kwargs)1125 # In this case, all layers are chunked1126 elif getattr(effective_config, "attention_chunk_size", None) is not None:1127 return create_chunked_causal_mask(**mask_kwargs)1128 # All layers use standard causal attention1129 return create_causal_mask(**mask_kwargs)1130 1131 1132# Below are utilities to pretty-print the different masks1133# Print the matrix with words as row labels1134GREEN = "\033[92m"1135YELLOW = "\033[93m"1136RESET = "\033[0m"1137BLACK_SQUARE = "■"1138WHITE_SQUARE = "⬚"1139GREY_SQUARE = "∙"1140LOW_TRIANGLE = "⬕"1141UPPER_TRIANGLE = "⬔"1142 1143 1144def get_style(style):1145 if style == "majong":1146 BLACK_SQUARE = "🀞" # Full block (represents "on" or active)1147 BLACK_SQUARE = "🀙" # Full block (represents "on" or active)1148 WHITE_SQUARE = "🀆" # "▒" # Light shade (represents "off" or inactive)1149 LOW_TRIANGLE = "🀛" # Lower left triangle (stylized indication)1150 UPPER_TRIANGLE = "🀛" # Upper left triangle (stylized indication)1151 else:1152 BLACK_SQUARE = "█" # Full block (represents "on" or active)1153 WHITE_SQUARE = "░" # "▒" # Light shade (represents "off" or inactive)1154 LOW_TRIANGLE = "▙" # Lower left triangle (stylized indication))1155 UPPER_TRIANGLE = "▜" # Upper left triangle (stylized indication)1156 1157 return BLACK_SQUARE, WHITE_SQUARE, LOW_TRIANGLE, UPPER_TRIANGLE1158 1159 1160# LOW_TRIANGLE = UPPER_TRIANGLE = "⟍" # Upper right triangle (stylized indication)1161 1162YELLOW_SQUARE = f"{YELLOW}{BLACK_SQUARE}{RESET}"1163GREEN_SQUARE = f"{GREEN}{BLACK_SQUARE}{RESET}"1164 1165 1166def tensor_to_mask_visual(original_tensor: torch.Tensor, grid_size=(20, 40), style="majong") -> str:1167 BLACK_SQUARE, WHITE_SQUARE, LOW_TRIANGLE, UPPER_TRIANGLE = get_style(style)1168 h, w = original_tensor.shape1169 max_h, max_w = grid_size1170 if not (h < max_h and w < max_w):1171 # Preserve aspect ratio within max grid size1172 aspect_ratio = 2 * w / h1173 if aspect_ratio > 1:1174 w = max_w1175 h = min(max_h, max(1, round(max_w / aspect_ratio)))1176 else:1177 h = max_h1178 w = max(1, round(max_h * aspect_ratio))1179 1180 # Step 1: Rescale tensor by average pooling1181 tensor = original_tensor.unsqueeze(0).unsqueeze(0) # Add batch and channel dimensions1182 tensor = F.adaptive_avg_pool2d(tensor, output_size=(h, w))[0, 0] # Remove extra dims1183 else:1184 tensor = original_tensor1185 1186 # Step 3: Build the string representation1187 result = []1188 for i in range(h):1189 row = ""1190 for j in range(w):1191 if tensor[i, j] == 1:1192 row += BLACK_SQUARE1193 elif tensor[i, j] == 0:1194 row += WHITE_SQUARE1195 else:1196 if j > 0:1197 if tensor[i, j - 1] == 1:1198 row += LOW_TRIANGLE1199 elif tensor[i, j - 1] == 0:1200 row += UPPER_TRIANGLE