Style

Author: remi-or

examples/pytorch/continuous_batching.py

@@ -241,7 +241,9 @@ def batch_generate(
     if args.output_file is None:
         os.makedirs("runs/cb", exist_ok=True)
         attn = args.attn.replace("|", "_").replace("/", "_")
-        args.output_file = f"runs/cb/{args.num_blocks}_{args.max_batch_tokens}_{attn}_{args.matmul_precision}_{args.samples}.json"
+        args.output_file = (
+            f"runs/cb/{args.num_blocks}_{args.max_batch_tokens}_{attn}_{args.matmul_precision}_{args.samples}.json"
+        )
 
     # Run warmup batch generation
     batch_generate(

examples/pytorch/continuous_batching_simple.py

@@ -48,14 +48,14 @@
     # Prepare tokenizer and dataset
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
     dataset = datasets.load_dataset("openai/gsm8k", "socratic", split="test")
-    dataset = dataset.select(range(args.samples)) 
+    dataset = dataset.select(range(args.samples))
     tokenized_datasets = dataset.map(lambda x: tokenizer(x["question"]), batched=True)
     simple_batch_inputs = [item["input_ids"] for item in tokenized_datasets]
 
     # Prepare generation config
     generation_config = GenerationConfig(
         max_new_tokens=512,
-        use_cuda_graph=False, # Not supported for simple version
+        use_cuda_graph=False,  # Not supported for simple version
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.pad_token_id,
         do_sample=False,

src/transformers/generation/continuous_batching/__init__.py

@@ -13,8 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from .cache import PagedAttentionCache
-from .continuous_api import ContinuousBatchingManager, ContinuousMixin
 from .classes import RequestState, RequestStatus
+from .continuous_api import ContinuousBatchingManager, ContinuousMixin
 
 
 __all__ = ["PagedAttentionCache", "RequestState", "RequestStatus", "ContinuousMixin", "ContinuousBatchingManager"]

src/transformers/generation/continuous_batching/cache.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 from collections import deque
 from math import floor, sqrt
-from typing import Any, Optional, TypeVar, Union
+from typing import Optional, Union
 
 import torch
 
@@ -287,9 +287,7 @@ def compute_num_blocks_and_max_batch_tokens(
         logger.info(f"Cache memory: {cache_memory}")
 
         # Compute memory footprints
-        mem_per_activation_token = (
-            m * self._activation_dtype.itemsize * (self.hidden_size + self.vocab_size)
-        )
+        mem_per_activation_token = m * self._activation_dtype.itemsize * (self.hidden_size + self.vocab_size)
         mem_per_cache_token = 2 * self.num_heads * self.head_dim * self.num_layers * cache_dtype.itemsize
         mem_per_input_token = 8 * m * self._input_dtype.itemsize
         logger.info(f"Memory per activation token: {mem_per_activation_token}")
@@ -299,7 +297,7 @@ def compute_num_blocks_and_max_batch_tokens(
         # Compute second-degree polynomial coefficients
         a = m * self._activation_dtype.itemsize
         b = mem_per_input_token + mem_per_cache_token + mem_per_activation_token
-        c = - cache_memory
+        c = -cache_memory
 
         # Compute discriminant and greatest solution
         discriminant = b**2 - 4 * a * c

src/transformers/generation/continuous_batching/continuous_api.py

@@ -106,7 +106,7 @@ def __init__(
         self.decode_stream = DecodeStream(skip_special_tokens=True)
 
     def return_attention_mask(self) -> bool:
-        return self.config._attn_implementation != "paged_attention" # we set `is_causal` to True in paged call
+        return self.config._attn_implementation != "paged_attention"  # we set `is_causal` to True in paged call
 
     @traced(standalone=True)
     def setup_static_tensors(self):
@@ -161,7 +161,6 @@ def reset_static_tensors(self):
         if self.attention_mask is not None:
             self.attention_mask[:, :, :t, :c].fill_(torch.finfo(self.model_dtype).min)
 
-
     def get_model_kwargs(self) -> PagedAttentionArgs:
         """Get model keyword arguments for the current batch."""
         # Compute the slice to return
@@ -172,8 +171,8 @@ def get_model_kwargs(self) -> PagedAttentionArgs:
             "input_ids": self.input_ids[:, :t],
             "attention_mask": self.attention_mask,
             "position_ids": self.position_ids[:, :t],
-            "cu_seq_lens_q": self.cumulative_seqlens_q[:t+1],
-            "cu_seq_lens_k": self.cumulative_seqlens_k[:t+1],
+            "cu_seq_lens_q": self.cumulative_seqlens_q[: t + 1],
+            "cu_seq_lens_k": self.cumulative_seqlens_k[: t + 1],
             "write_index": self.write_index[:t],
             "read_index": self.read_index[:c],
             "logits_indices": self.logits_indices[:t],