Merge pull request #457 from datalab-to/dev

Move flash attention funcs

Author: VikParuchuri

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "surya-ocr"
-version = "0.16.6"
+version = "0.16.7"
 description = "OCR, layout, reading order, and table recognition in 90+ languages"
 authors = ["Vik Paruchuri <[email protected]>"]
 readme = "README.md"

surya/common/surya/__init__.py

@@ -17,13 +17,8 @@
 from surya.common.surya.encoder import SuryaEncoderModel
 from surya.settings import settings
 
-from transformers.utils import is_flash_attn_2_available
-
 from surya.logging import get_logger
 
-if is_flash_attn_2_available():
-    from surya.common.surya.flash_attn_utils import _get_unpad_data
-
 logger = get_logger()
 
 
@@ -415,6 +410,9 @@ def forward(
         # Handling flash attention kwargs outside the decoder to speed up + avoid graph breaks inside the decoder
         # Skipped during decoding since not required
         if self.decoder.config._attn_implementation == "flash_attention_2" and prefill:
+            # Needed for CPU -> GPU
+            from surya.common.surya.flash_attn_utils import _get_unpad_data
+
             batch_size, query_length, _ = inputs_embeds.shape
             indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(
                 attention_mask

surya/common/surya/decoder/__init__.py

@@ -21,13 +21,6 @@
 from surya.common.pretrained import SuryaPreTrainedModel
 from surya.common.surya.decoder.config import SuryaDecoderConfig
 
-from transformers.utils import is_flash_attn_2_available
-
-if is_flash_attn_2_available():
-    from surya.common.surya.flash_attn_utils import (
-        flash_attn_decode,
-        flash_attn_prefill,
-    )
 
 logger = logging.get_logger(__name__)
 
@@ -206,6 +199,12 @@ def forward(
                     'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
                 )
             elif self.config._attn_implementation == "flash_attention_2":
+                # Needed for CPU -> GPU
+                from surya.common.surya.flash_attn_utils import (
+                    flash_attn_decode,
+                    flash_attn_prefill,
+                )
+
                 if prefill:
                     attention_interface = flash_attn_prefill
                 else:

surya/common/surya/encoder/__init__.py

@@ -5,17 +5,11 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from transformers.activations import ACT2FN
-from transformers.utils import is_flash_attn_2_available
 
 from surya.common.pretrained import SuryaPreTrainedModel
 from surya.common.surya.encoder.config import SuryaEncoderConfig
 from surya.logging import get_logger
 
-if is_flash_attn_2_available():
-    from flash_attn import flash_attn_varlen_func
-    from flash_attn.layers.rotary import apply_rotary_emb  # noqa
-
-
 logger = get_logger()
 
 
@@ -127,6 +121,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 def apply_rotary_pos_emb_flashatt(
     q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
 ) -> Tuple[torch.Tensor, torch.Tensor]:
+    from flash_attn.layers.rotary import apply_rotary_emb
+
     cos = cos.chunk(2, dim=-1)[0].contiguous()
     sin = sin.chunk(2, dim=-1)[0].contiguous()
     q_embed = apply_rotary_emb(q.float(), cos.float(), sin.float()).type_as(q)
@@ -148,6 +144,8 @@ def forward(
         rotary_pos_emb: Optional[torch.Tensor] = None,
         position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> torch.Tensor:
+        from flash_attn import flash_attn_varlen_func
+
         bsz = hidden_states.shape[0]
         seq_length = hidden_states.shape[1]
         q, k, v = (

surya/ocr_error/model/encoder.py

@@ -16,16 +16,11 @@
 )
 
 from transformers.utils import (
-    is_flash_attn_2_available,
     is_flash_attn_greater_or_equal_2_10,
 )
 
 from surya.common.pretrained import SuryaPreTrainedModel
 
-if is_flash_attn_2_available():
-    from flash_attn import flash_attn_func, flash_attn_varlen_func
-    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa
-
 from surya.common.s3 import S3DownloaderMixin
 from surya.ocr_error.model.config import DistilBertConfig
 
@@ -342,6 +337,9 @@ def _flash_attention_forward(
             softmax_scale (`float`, *optional*):
                 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
         """
+        from flash_attn import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.bert_padding import pad_input
+
         if not self._flash_attn_uses_top_left_mask:
             causal = self.is_causal
         else:
@@ -397,6 +395,8 @@ def _flash_attention_forward(
     def _upad_input(
         self, query_layer, key_layer, value_layer, attention_mask, query_length
     ):
+        from flash_attn.bert_padding import index_first_axis, unpad_input
+
         indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
         batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape