Aluode/PerceptionLabPortable
0
1"""2Partially inspired by torchtune's flex attention implementation3 4Citation:5@software{torchtune,6 title = {torchtune: PyTorch's finetuning library},7 author = {torchtune maintainers and contributors},8 url = {https//github.com/pytorch/torchtune},9 license = {BSD-3-Clause},10 month = apr,11 year = {2024}12}13"""14# coding=utf-815# Copyright 2025 The HuggingFace Inc. team.16#17# Licensed under the Apache License, Version 2.0 (the "License");18# you may not use this file except in compliance with the License.19# You may obtain a copy of the License at20#21# http://www.apache.org/licenses/LICENSE-2.022#23# Unless required by applicable law or agreed to in writing, software24# distributed under the License is distributed on an "AS IS" BASIS,25# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.26# See the License for the specific language governing permissions and27# limitations under the License.28 29from typing import Optional, Union30 31import torch32from packaging import version33 34from ..utils import is_torch_flex_attn_available, logging35from ..utils.import_utils import _torch_version, is_torch_less_or_equal, is_torchdynamo_compiling36 37 38if is_torch_flex_attn_available():39 from torch.nn.attention.flex_attention import _DEFAULT_SPARSE_BLOCK_SIZE as flex_default_block_size40 from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention41 42 43logger = logging.get_logger(__name__)44 45 46class WrappedFlexAttention:47 """48 We are doing a singleton class so that flex attention is compiled once when it's first called.49 """50 51 _instance = None52 _is_flex_compiled = False53 _compiled_flex_attention = None54 55 def __new__(cls, *args, **kwargs):56 if cls._instance is None:57 # Create a new instance if one doesn't already exist58 cls._instance = super().__new__(cls)59 return cls._instance60 61 @torch.compiler.disable(recursive=False)62 def __init__(self, training):63 """64 Initialize or update the singleton instance.65 """66 if not self._is_flex_compiled or training != self.training:67 self.training = training68 if is_torch_less_or_equal("2.5.1"):69 self._compiled_flex_attention = torch.compile(flex_attention, dynamic=False)70 # In PyTorch 2.6.0, there's a known issue with flex attention compilation which may71 # cause errors. The suggested fix is to compile with "max-autotune-no-cudagraphs"72 # see https://github.com/pytorch/pytorch/issues/146260 for training73 elif version.parse(_torch_version).base_version == "2.6.0" and training:74 self._compiled_flex_attention = torch.compile(75 flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"76 )77 # Fallback, usually the most recent torch 2.7.x+ versions78 else:79 self._compiled_flex_attention = torch.compile(flex_attention)80 81 self._is_flex_compiled = True82 83 def __call__(self):84 return self._compiled_flex_attention85 86 87def compile_friendly_flex_attention(88 query: torch.Tensor,89 key: torch.Tensor,90 value: torch.Tensor,91 training=False,92 **kwargs,93) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:94 # First call initialise singleton wrapper object, second call invokes the object method to return compiled flex attention95 # Do not use compiled version if already compiling forward (it raises issues)96 flex_attention_compiled = WrappedFlexAttention(training)() if not is_torchdynamo_compiling() else flex_attention97 return flex_attention_compiled(98 query,99 key,100 value,101 **kwargs,102 )103 104 105Offset = Union[torch.Tensor, int]106 107 108# TODO: deprecate / rename to make_flex_block_mask for clarity as it's not only causal anymore109def make_flex_block_causal_mask(110 attention_mask_2d: torch.Tensor,111 attention_chunk_size: Optional[int] = None,112 query_length=None,113 key_length=None,114 offsets: Optional[tuple[Offset, Offset]] = None,115 is_causal: Optional[bool] = True,116) -> "BlockMask":117 """118 IMPORTANT NOTICE: This function is deprecated in favor of using the mask primitives in `masking_utils.py`,119 and will be removed in a future version without warnings. New code should not use it. It is only kept here120 for BC for now, while models using it are being patched accordingly.121 122 Create a block (causal) document mask for a batch of sequences, both packed and unpacked.123 Create Block (causal) logic and passing it into :func:`torch.nn.attention.flex_attention.create_block_mask`.124 The resultant BlockMask is a compressed representation of the full (causal) block125 mask. BlockMask is essential for performant computation of flex attention.126 See: https://pytorch.org/blog/flexattention/127 128 Args:129 attention_mask_2d (torch.Tensor): Attention mask for packed and padded sequences130 of shape (batch_size, total_seq_len). e.g.131 132 For unpacked sequence:133 [[1, 1, 1, 1, 0, 0, 0],134 [1, 1, 1, 1, 1, 0, 0]]135 136 For packed sequence:137 [[1, 1, 1, 2, 2, 2, 0],138 [1, 1, 2, 2, 2, 3, 3]]139 140 Returns:141 BlockMask142 """143 batch_size, total_seq_len = attention_mask_2d.shape144 if not key_length:145 key_length = total_seq_len146 if not query_length:147 query_length = total_seq_len148 # older torch (2.5.x) cannot handle sequences not in multiples of 128 (default block size)149 pad_len = ((key_length // flex_default_block_size) + 1) * flex_default_block_size150 attention_mask_2d = torch.nn.functional.pad(attention_mask_2d, value=0, pad=(0, pad_len - key_length))151 device = attention_mask_2d.device152 document_ids = attention_mask_2d.clone()153 154 if attention_chunk_size is not None:155 # we create an arange, then we just // by chunk size to get [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3]156 chunk_idxs = (document_ids.clone().fill_(1).cumsum(-1) - 1) // (attention_chunk_size)157 158 # Instead of passing a tensor mask, flex attention requires a mask_mod function159 # that determines which elements of QK^T should be included in the attention160 # computation prior to the softmax. For sample packing, we need both the161 # logic for both causal mask and document mask. See PyTorch's official162 # blog post for more details: https://pytorch.org/blog/flexattention/#mask-mods163 def causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):164 """165 Defines the logic of a block causal mask by combining both a standard causal mask166 and a block diagonal document mask.167 See :func:`~torchtune.modules.attention_utils.create_block_causal_mask`168 for an illustration.169 """170 causal_mask = q_idx >= kv_idx # not valid when decoding171 document_mask = document_ids[batch_idx, q_idx] == document_ids[batch_idx, kv_idx]172 padding_mask = attention_mask_2d[batch_idx, q_idx] > 0173 final_mask = causal_mask & padding_mask & document_mask174 return final_mask175 176 def chunk_causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx):177 """178 Combines the chunk mask with the causal mask for chunked attention.179 """180 chunk_mask = chunk_idxs[batch_idx, q_idx] == chunk_idxs[batch_idx, kv_idx]181 causal_doc_mask = causal_mask_mod(batch_idx, head_idx, q_idx, kv_idx)182 return chunk_mask & causal_doc_mask183 184 def default_mask_mod(batch_idx, head_idx, q_idx, kv_idx):185 """186 Utilizes default attention mask to enable encoder and encoder-decoder187 attention masks.188 """189 document_mask = document_ids[batch_idx, q_idx] == document_ids[batch_idx, kv_idx]190 # kv indexing is crucial in order to work correctly191 padding_mask = attention_mask_2d[batch_idx, kv_idx] > 0192 final_mask = padding_mask & document_mask193 return final_mask194 195 if not is_causal:196 mask_mod_maybe_combined = default_mask_mod197 else:198 mask_mod_maybe_combined = causal_mask_mod if attention_chunk_size is None else chunk_causal_mask_mod199 200 if offsets is not None:201 q_offset = offsets[0].to(device)202 kv_offset = offsets[1].to(device)203 204 def mask_mod(batch_idx, head_idx, q_idx, kv_idx):205 offset_q = q_idx + q_offset206 offset_kv = kv_idx + kv_offset207 return mask_mod_maybe_combined(batch_idx, head_idx, offset_q, offset_kv)208 else:209 mask_mod = mask_mod_maybe_combined210 211 return create_block_mask(212 mask_mod=mask_mod,213 B=batch_size,214 H=None, # attention head215 Q_LEN=query_length,216 KV_LEN=key_length,217 device=device,218 # compiling the mask is not BC with older torch219 _compile=not is_torch_less_or_equal("2.5.1"),220 )221 222 223def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:224 """225 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,226 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)227 """228 batch, num_key_value_heads, slen, head_dim = hidden_states.shape229 if n_rep == 1:230 return hidden_states231 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)232 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)233 234 235def flex_attention_forward(236 module: torch.nn.Module,237 query: torch.Tensor,238 key: torch.Tensor,239 value: torch.Tensor,240 attention_mask: Union[torch.Tensor, "BlockMask"],241 scaling: Optional[float] = None,242 softcap: Optional[float] = None,243 head_mask: Optional[torch.Tensor] = None,244 s_aux: Optional[torch.Tensor] = None,245 **kwargs,246) -> tuple[torch.Tensor, Optional[torch.Tensor]]:247 if head_mask is not None:248 logger.warning_once(249 "`flex_attention` does not support `head_mask`. Please set your attention to `eager` if you want this feature."250 )251 252 if kwargs.get("dropout", 0.0) > 0:253 raise ValueError(254 "`flex_attention` does not support `dropout`. Please use it with inference"255 " only (`model.eval()`) or turn off the attention dropout in the respective config."256 )257 258 block_mask = None259 score_mask = None260 if isinstance(attention_mask, BlockMask):261 block_mask = attention_mask262 else:263 score_mask = attention_mask264 265 if score_mask is not None:266 score_mask = score_mask[:, :, :, : key.shape[-2]]267 268 def score_mod(score, batch_idx, head_idx, q_idx, kv_idx):269 if softcap is not None:270 score = softcap * torch.tanh(score / softcap)271 if score_mask is not None:272 score = score + score_mask[batch_idx][0][q_idx][kv_idx]273 if head_mask is not None:274 score = score + head_mask[batch_idx][head_idx][0][0]275 # Note: attention sinks cannot be correctly implemented in score_mod276 # because it requires operating on the full attention matrix before softmax.277 # ==> this is done after flex attention278 return score279 280 enable_gqa = True281 num_local_query_heads = query.shape[1]282 283 # When running TP this helps:284 if (num_local_query_heads & (num_local_query_heads - 1)) != 0:285 key = repeat_kv(key, query.shape[1] // key.shape[1])286 value = repeat_kv(value, query.shape[1] // value.shape[1])287 enable_gqa = False288 289 kernel_options = kwargs.get("kernel_options")290 # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it291 return_lse = query.device.type != "cpu"292 293 if not return_lse and s_aux is not None:294 raise ValueError(295 "Attention sinks cannot be run on CPU with flex attention. Please switch to a different device, e.g. CUDA"296 )297 298 flex_attention_output = compile_friendly_flex_attention(299 query,300 key,301 value,302 score_mod=score_mod,303 block_mask=block_mask,304 enable_gqa=enable_gqa,305 scale=scaling,306 kernel_options=kernel_options,307 # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.308 # For simplification, we thus always return it as no additional computations are introduced.309 return_lse=return_lse,310 training=module.training,311 )312 # lse is returned in float32313 if return_lse:314 attention_output, lse = flex_attention_output # type: ignore[misc]315 lse = lse.to(value.dtype)316 317 if s_aux is not None:318 # Apply attention sinks by renormalizing using LSE319 batch_size, num_heads, seq_len_q, _ = attention_output.shape # batch, num_heads, seq_len, head_dim320 sinks = s_aux.view(1, -1, 1, 1).expand(batch_size, num_heads, seq_len_q, 1)321 322 # We need to compute the normalization that includes the sinks323 # since log(sum(exp(scores))) = lse, exp(log(sum(exp(scores)))) = exp(lse)324 # NB: log(sum(exp(scores)) + exp(sink)) = log(exp(lse) + exp(sink))325 lse_expanded = lse.unsqueeze(-1) # [batch, num_heads, seq_len, 1]326 combined_lse = torch.logsumexp(torch.cat([lse_expanded, sinks], dim=-1), dim=-1, keepdim=True)327 328 # Use new_norm / old_norm = exp(combined_lse - lse) to compute renorm and apply329 renorm_factor = torch.exp(lse_expanded - combined_lse)330 attention_output = attention_output * renorm_factor331 else:332 attention_output = flex_attention_output # type: ignore[assignment]333 lse = None334 335 attention_output = attention_output.transpose(1, 2).contiguous()336 return attention_output, lse337