transformers 5.0.0rc3__py3-none-any.whl → 5.1.0__py3-none-any.whl

transformers/file_utils.py

@@ -63,10 +63,8 @@ from .utils import (
     is_datasets_available,
     is_detectron2_available,
     is_faiss_available,
-    is_ftfy_available,
     is_g2p_en_available,
     is_in_notebook,
-    is_ipex_available,
     is_librosa_available,
     is_onnx_available,
     is_pandas_available,

transformers/generation/candidate_generator.py

@@ -27,7 +27,6 @@ from ..utils import is_sklearn_available
 if is_sklearn_available():
     from sklearn.metrics import roc_curve
 
-from ..pytorch_utils import isin_mps_friendly
 from .logits_process import LogitsProcessorList, MinLengthLogitsProcessor, SuppressTokensLogitsProcessor
 
 
@@ -1165,7 +1164,7 @@ class PromptLookupCandidateGenerator(CandidateGenerator):
                     # remove remaining candidate ids if an "eos" token is found, otherwise the target model may
                     # accept eos and the rest as valid, thus not stopping generation after "eos"
                     # NOTE: below code is written based on the fact that assisted decoding supports only bs=1
-                    mask = isin_mps_friendly(chosen_ids, self.eos_token_id)
+                    mask = torch.isin(chosen_ids, self.eos_token_id)
                     match_indices_eos = torch.nonzero(mask)
                     if match_indices_eos.numel() > 0:
                         first_eos_index = match_indices_eos[0].item()

transformers/generation/continuous_batching/cache.py

@@ -17,6 +17,7 @@ import torch
 
 from ...configuration_utils import PreTrainedConfig
 from ...generation.configuration_utils import GenerationConfig
+from ...utils.generic import is_flash_attention_requested
 from ...utils.metrics import attach_tracer, traced
 from .cache_manager import BlockManager, CacheAllocator, FullAttentionCacheAllocator, SlidingAttentionCacheAllocator
 from .requests import RequestState, get_device_and_memory_breakdown, logger
@@ -117,7 +118,7 @@ class PagedAttentionCache:
         self,
         config: PreTrainedConfig,
         generation_config: GenerationConfig,
-        device: torch.device,
+        device: torch.device | str,
         dtype: torch.dtype = torch.float16,
         tp_size: int | None = None,
         allow_block_sharing: bool = True,
@@ -172,7 +173,7 @@ class PagedAttentionCache:
         # Infer number of blocks and max batch tokens
         page_size = self.head_dim * self.num_key_value_heads
 
-        if "flash" in self.config._attn_implementation:
+        if is_flash_attention_requested(self.config):
             num_attention_masks = 0  # only used to compute the default memory footprint args
         elif "sliding_attention" in group_types:
             # TODO: when we generalize to allow for block-attn, we can use `num_attention_masks=sum(set(group_types))`
@@ -200,6 +201,7 @@ class PagedAttentionCache:
         # Add the inferred attributes to the class
         self.num_blocks = num_blocks
         self.max_batch_tokens = max_batch_tokens
+        self.num_pages = self.num_blocks * self.block_size
         logger.info(
             f"PagedAttentionCache initialized with {self.num_blocks = }, {self.block_size = }, {page_size = }, "
             f"{self.max_batch_tokens = } {num_attention_masks = }"
@@ -290,35 +292,33 @@ class PagedAttentionCache:
         return self._block_manager.num_free_blocks
 
     @traced
-    def extend_read_indices(
-        self, request_id: str, past_length: int, query_length: int, read_index: list[list[int]]
+    def extend_read_and_write_indices(
+        self,
+        request_id: str,
+        past_length: int,
+        query_length: int,
+        read_index: list[list[int]],
+        write_index: list[list[int]],
     ) -> None:
         """Retrieve physical cache indices for reading KV states in the cache across all layer groups. This method
         coordinates with all cache managers to build the complete set of read indices needed for attention computation.
         """
-        for cm, read_indices in zip(self.group_cache_managers, read_index):
+        for cm, read_indices, write_indices in zip(self.group_cache_managers, read_index, write_index):
             indices = cm.get_read_indices(request_id, past_length, query_length)
             read_indices.extend(indices)
-
-    @traced
-    def extend_write_indices(
-        self, request_id: str, past_length: int, query_length: int, write_index: list[list[int]]
-    ) -> None:
-        """Retrieve physical cache indices for writing new KV states to the cache across all layer groups. This method
-        coordinates with all cache managers to build the complete set of write indices needed to store computed KV
-        states."""
-        for cm, write_indices in zip(self.group_cache_managers, write_index):
             indices = cm.get_write_indices(request_id, past_length, query_length)
             write_indices.extend(indices)
 
     @traced
-    def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> dict[str, int]:
+    def get_seqlens_k(self, past_length: int, query_length: int) -> dict[str, int]:
         """Retrieve the key sequence length for the given request_id across all layer types. Returns a dictionary of
         layer types to their corresponding key sequence lengths."""
         seqlens_k = {}
-        for cm in self.group_cache_managers:
-            attn_type, seqlen_k = cm.get_seqlens_k(request_id, past_length, query_length)
-            seqlens_k[attn_type] = seqlen_k
+        if self.num_full_attention_groups > 0:
+            seqlens_k["full_attention"] = past_length + query_length
+        if self.num_sliding_attention_groups > 0:
+            seqlens_k["sliding_attention"] = query_length + min(past_length, self.config.sliding_window - 1)
+        # NOTE: when we add more attention types / different sliding windows, we can go back to looping over CMs
         return seqlens_k
 
     @traced
@@ -416,10 +416,10 @@ class PagedAttentionCache:
                     prompt_ids=(state.initial_tokens + state.generated_tokens),
                 )
 
-    def copy_cache(self, source_blocks: list[int], forked_blocks: list[int]) -> None:
+    def copy_cache(self, list_source_blocks: list[int], list_forked_blocks: list[int]) -> None:
         """Copy the cache from the source blocks to the forked blocks."""
-        source_blocks = torch.tensor(source_blocks, device=self.device, dtype=torch.int32)
-        forked_blocks = torch.tensor(forked_blocks, device=self.device, dtype=torch.int32)
+        source_blocks = torch.tensor(list_source_blocks, device=self.device, dtype=torch.int32)
+        forked_blocks = torch.tensor(list_forked_blocks, device=self.device, dtype=torch.int32)
         for key_cache, value_cache in zip(self.key_cache, self.value_cache):
             key_cache = key_cache.view(-1, self.block_size, self.num_key_value_heads, self.head_dim)
             value_cache = value_cache.view(-1, self.block_size, self.num_key_value_heads, self.head_dim)
@@ -533,24 +533,26 @@ class PagedAttentionMemoryHandler:
 
         where we already simplified int32_size = 4.
         """
-        # If neither num_blocks nor max_batch_tokens are provided, we use a second-order polynomial
-        if num_blocks is None and max_batch_tokens is None:
-            num_blocks, max_batch_tokens = self.compute_num_blocks_and_max_batch_tokens(
-                max_memory_percent, cache_dtype
-            )
-        # If only num_blocks is provided, we infer the max_batch_tokens
-        elif num_blocks is not None and max_batch_tokens is None:
+        if num_blocks is None:
+            if max_batch_tokens is None:
+                # If neither num_blocks nor max_batch_tokens are provided, we use a second-order polynomial
+                num_blocks, max_batch_tokens = self.compute_num_blocks_and_max_batch_tokens(
+                    max_memory_percent, cache_dtype
+                )
+            else:
+                # If only max_batch_tokens is provided, we infer the num_blocks
+                num_blocks = self.compute_num_blocks(max_batch_tokens, max_memory_percent, cache_dtype)
+        elif max_batch_tokens is None:
+            # If only num_blocks is provided, we infer the max_batch_tokens
             max_batch_tokens = self.compute_max_batch_tokens(num_blocks, max_memory_percent, cache_dtype)
-        # If only max_batch_tokens is provided, we infer the num_blocks
-        elif max_batch_tokens is not None and num_blocks is None:
-            num_blocks = self.compute_num_blocks(max_batch_tokens, max_memory_percent, cache_dtype)
+        else:
+            # If both num_blocks and max_batch_tokens are provided, we use them (useless, but helps with typing)
+            max_batch_tokens = max_batch_tokens
 
         # We check if the memory footprint is too large in all cases
         available_memory = self.get_available_memory(max_memory_percent)
         memory_footprint = self.compute_memory_footprint(
-            max_batch_tokens=max_batch_tokens,
-            num_blocks=num_blocks,
-            cache_dtype=cache_dtype,
+            max_batch_tokens=max_batch_tokens, num_blocks=num_blocks, cache_dtype=cache_dtype
         )
         if memory_footprint > available_memory:
             raise MemoryError(f"Memory footprint {memory_footprint} is more than available memory {available_memory}")
@@ -669,10 +671,10 @@ class PagedAttentionMemoryHandler:
 
     def compute_memory_footprint(
         self,
-        num_blocks: int | None = None,
-        max_batch_tokens: int | None = None,
-        cache_dtype: torch.dtype = torch.float16,
-    ) -> tuple[int, int, int]:
+        num_blocks: int,
+        max_batch_tokens: int,
+        cache_dtype: torch.dtype,
+    ) -> int:
         """Calculate the memory footprint breakdown for a given number of blocks and maximum batch tokens. The memory
         footprint is given by:
 

transformers/generation/continuous_batching/cache_manager.py

@@ -99,7 +99,8 @@ class BlockManager:
         for _ in range(block_to_uninitialize):
             id_to_uninitialize = self._init_block_ids.popitem()[0]
             block = self._id_to_block[id_to_uninitialize]
-            self._hash_to_id.pop(block.hash)
+            # Since the block is initialized it must have a hash, thus no need to check .hash is not None
+            self._hash_to_id.pop(block.hash)  # ty:ignore[invalid-argument-type]
             self._uninit_block_ids.append(id_to_uninitialize)
         return True
 
@@ -124,7 +125,7 @@ class BlockManager:
 
     def fork_blocks(
         self, parent_blocks: list[int], num_forks: int, shareable: bool, group_id: int
-    ) -> tuple[list[list[int]], list[int], list[int]]:
+    ) -> tuple[list[list[int]] | None, list[int], list[int]]:
         """Fork a given list of (parent_blocks) as many times as (num_forks). If the blocks are (shareable), we use
         reference on the blocks that are complete. Otherwise, we allocate new blocks and keep track of their indices to
         later copy the physical cache. For instance, when forking 4 blocks for 2 children:
@@ -307,10 +308,6 @@ class CacheAllocator(ABC):
     def get_write_indices(self, request_id: str, past_length: int, query_length: int) -> list[int]:
         """Returns the physical indices of where to write request_id's cache in the cache tensor."""
 
-    @abstractmethod
-    def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> tuple[str, int]:
-        """Returns the attention type of the cache allocator and the key sequence length for the given request_id."""
-
     def fork_blocks(
         self, parent_request_id: str, children_request_ids: list[str], block_manager: BlockManager
     ) -> tuple[list[int], list[int]]:
@@ -405,11 +402,6 @@ class FullAttentionCacheAllocator(CacheAllocator):
             physical_indices.append(physical_index)
         return physical_indices
 
-    def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> tuple[str, int]:
-        """Returns the attention type of the cache allocator and the key sequence length for the given request_id."""
-        seqlens_k = past_length + query_length
-        return "full_attention", seqlens_k
-
 
 class SlidingAttentionCacheAllocator(CacheAllocator):
     """Cache manager for sliding window attention layers."""
@@ -496,8 +488,3 @@ class SlidingAttentionCacheAllocator(CacheAllocator):
         if padding_length > 0:
             physical_indices = [-1] * padding_length + physical_indices
         return physical_indices
-
-    def get_seqlens_k(self, request_id: str, past_length: int, query_length: int) -> tuple[str, int]:
-        """Returns the attention type of the cache allocator and the key sequence length for the given request_id."""
-        seqlens_k = query_length + min(past_length, self.sliding_window - 1)
-        return "sliding_attention", seqlens_k