Repository: huggingface/transformers
Pull request: enable cpu paged cache

CPU can also use paged cache with eager or sdpa:
`python continuous_batching_simple.py --attn sdpa`

Without this change, the previous command error would be like:
```
Error in generation loop: unsupported operand type(s) for -: 'NoneType' and 'int'
Traceback (most recent call last):
  File "/home/jiqing/transformers/src/transformers/generation/continuous_batching/continuous_api.py", line 1017, in _run_generation_loop
    paged_attention_cache = PagedAttentionCache(
                            ^^^^^^^^^^^^^^^^^^^^
  File "/home/jiqing/transformers/src/transformers/generation/continuous_batching/cache.py", line 191, in __init__
    num_blocks, max_batch_tokens = memory_handler.infer_num_blocks_and_max_batch_tokens(
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jiqing/transformers/src/transformers/generation/continuous_batching/cache.py", line 481, in infer_num_blocks_and_max_batch_tokens
    num_blocks, max_batch_tokens = self.compute_num_blocks_and_max_batch_tokens(
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jiqing/transformers/src/transformers/generation/continuous_batching/cache.py", line 522, in compute_num_blocks_and_max_batch_tokens
    cache_memory = self.get_available_memory(max_memory_percent)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jiqing/transformers/src/transformers/generation/continuous_batching/cache.py", line 456, in get_available_memory
    available_memory = total - max(allocated, reserved)
                       ~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~
TypeError: unsupported operand type(s) for -: 'NoneType' and 'int'

```

Changed files:
filename: examples/pytorch/continuous_batching_simple.py
patch: @@ -19,6 +19,7 @@
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers.generation import GenerationConfig
+from transformers.utils import is_torch_accelerator_available
 
 
 MODEL_ID = "Qwen/Qwen3-4B-Instruct-2507"
@@ -36,11 +37,14 @@
 
     args = parser.parse_args()
 
+    device = torch.accelerator.current_accelerator() if is_torch_accelerator_available() else "cuda"
+    device_map = "cpu" if device is None else device.type
+
     # Prepare model
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         attn_implementation=args.attn,
-        device_map="cuda",
+        device_map=device_map,
         dtype=torch.bfloat16,
     )
     model = model.eval()

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -395,6 +395,7 @@ def get_model_kwargs(self, padded_q_size: int = 0, padded_kv_cache_size: int = 0
         # some NaNs in the output logits even for non-padded tokens.
         if use_padding:
             self.max_seqlen_q = max(self.max_seqlen_q, q_len - self.total_seqlen_q)
+            kwargs.max_seqlen_q = self.max_seqlen_q  # Update kwargs with the new value
             self.cumulative_seqlens_q[self.actual_batch_size + 1 :] = q_len
             # FIXME: is there another way to avoid this? It has a very slight impact on performance (~5 tok/s)
 
@@ -929,6 +930,11 @@ def _decide_use_cuda_graphs(
         If none of the above criteria are met, we use a default heuristic based on the attention implementation: we turn
         on cuda graphs if and only if no attention mask is needed.
         """
+        # If cuda is not available, we cannot use cuda graphs
+        if not torch.cuda.is_available():
+            if use_cuda_graph:
+                logger.warning(f"use_cuda_graph is True but {torch.cuda.is_available() = }: turning off cuda graphs.")
+            return False
         # If use_cuda_graph is specified, we follow the user's choice
         if use_cuda_graph is not None:
             return use_cuda_graph

filename: src/transformers/generation/continuous_batching/requests.py
patch: @@ -17,11 +17,14 @@
 
 import torch
 
-from ...utils import is_torch_xpu_available
+from ...utils import is_psutil_available, is_torch_xpu_available
 from ...utils.logging import logging
 from ...utils.metrics import traced
 
 
+if is_psutil_available():
+    import psutil
+
 # This is a temporary token ID used to represent a token that is not yet generated
 TMP_TOKEN_ID = -1
 
@@ -53,9 +56,19 @@ def get_device_and_memory_breakdown() -> tuple[torch.device, int, int, int]:
         reserved_memory = 0  # MPS does not track reserved separately
     else:
         device = torch.device("cpu")
-        total_memory = 0
-        reserved_memory = 0
-        allocated_memory = 0
+        if is_psutil_available():
+            total_memory = psutil.virtual_memory().total
+            allocated_memory = psutil.Process().memory_info().rss
+            reserved_memory = allocated_memory
+        else:
+            logger.error(
+                "Cannot get memory breakdown on CPU without psutil: returning 0 for all memory values. Please install "
+                "psutil to get an actual memory breakdown."
+            )
+            total_memory = 0
+            reserved_memory = 0
+            allocated_memory = 0
+
     return device, total_memory, reserved_memory, allocated_memory
 
 

filename: tests/generation/test_continuous_batching.py
patch: @@ -39,7 +39,6 @@
     Expectations,
     require_deterministic_for_xpu,
     require_flash_attn,
-    require_torch_accelerator,
     slow,
     torch_device,
 )
@@ -214,7 +213,6 @@ def test_attention_mask(
             (2, 0, 1, 1, 3, 4, True),
         ]
     )
-    @require_torch_accelerator
     def test_continuous_batching_will_allocation_be_successful(
         self,
         num_requested_blocks: int,
@@ -299,6 +297,13 @@ def _test_continuous_batching_parity(
 
         # Generation with continuous batching
         model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=attn_implementation, dtype=dtype)
+        if (
+            attn_implementation == "flash_attention_2"
+            and torch_device == "cpu"
+            and getattr(model.config, "sliding_window", None) is not None
+            and model.config.sliding_window > 0
+        ):
+            self.skipTest("Flash Attention 2 with sliding window attention is not supported on CPU. Skipping test.")
         model = model.to(torch_device).eval()
         model.generation_config.max_new_tokens = max_new_tokens
         model.generation_config.do_sample = False
@@ -372,7 +377,6 @@ def _test_continuous_batching_parity(
             )
         )
     )
-    @require_torch_accelerator
     @slow
     def test_continuous_batching_config_combinations(
         self,
@@ -398,25 +402,21 @@ def test_continuous_batching_config_combinations(
             )
         )
     )
-    @require_torch_accelerator
     @slow
     def test_continuous_batching_diverse_models(self, model_id: str, use_cuda_graph: bool, use_compile: bool) -> None:
         try:
             self._test_continuous_batching_parity(model_id, True, "flash_attention_2", use_cuda_graph, use_compile)
         finally:
             flush_memory(flush_compile=use_compile)
 
-    @require_torch_accelerator
     def test_continuous_batching_fast(self) -> None:
         model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         self._test_continuous_batching_parity(model_id, False, "sdpa", False, False)
 
-    @require_torch_accelerator
     def test_continuous_batching_long_generate(self) -> None:
         model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         self._test_continuous_batching_parity(model_id, True, "flash_attention_2", True, True, max_new_tokens=80)
 
-    @require_torch_accelerator
     def test_continuous_batching_few_blocks(self) -> None:
         """This test verifies that generation works with a very small number of blocks, ie. small enough that we need to
         offload a request at some point. To add more complexity, we repeat the same prompt 4 times and enable prefix
@@ -475,15 +475,12 @@ def _test_streaming_or_not_request(self, with_streaming: bool, with_non_streamin
 
         manager.stop(block=True)
 
-    @require_torch_accelerator
     def test_streaming_request(self) -> None:
         self._test_streaming_or_not_request(with_streaming=True, with_non_streaming=False)
 
-    @require_torch_accelerator
     def test_non_streaming_request(self) -> None:
         self._test_streaming_or_not_request(with_streaming=False, with_non_streaming=True)
 
-    @require_torch_accelerator
     def test_streaming_and_non_streaming_requests_can_alternate(self) -> None:
         self._test_streaming_or_not_request(with_streaming=True, with_non_streaming=True)
 
@@ -580,7 +577,6 @@ def _test_block_sharing(
         print(f"{chunk_no_reuse.generated_tokens = } {expected_output_tokens = }")
         self.assertEqual(chunk_no_reuse.generated_tokens, expected_output_tokens)
 
-    @require_torch_accelerator
     def test_prefix_sharing(self) -> None:
         model_id = "Qwen/Qwen2.5-0.5B-Instruct"
         num_layer_groups = {"full_attention": 1, "sliding_window": 0}
@@ -591,7 +587,6 @@ def test_prefix_sharing(self) -> None:
 
         return self._test_block_sharing(model_id, num_layer_groups, input_msg, expected_generated_tokens)
 
-    @require_torch_accelerator
     def test_block_sharing_with_hybrid_model(self) -> None:
         model_id = "google/gemma-3-1b-it"
         num_layer_groups = {"full_attention": 2, "sliding_window": 11}
@@ -603,7 +598,6 @@ def test_block_sharing_with_hybrid_model(self) -> None:
         return self._test_block_sharing(model_id, num_layer_groups, input_msg, expected_generated_tokens)
 
     @parameterized.expand([True, False])
-    @require_torch_accelerator
     @require_flash_attn  # otherwise the test can fail because attention bias has a very slight impact on SDPA and eager
     def test_num_return_sequences(self, allow_block_sharing: bool) -> None:
         model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

filename: tests/generation/test_paged_attention.py
patch: @@ -4,7 +4,7 @@
 from parameterized import parameterized
 
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
-from transformers.testing_utils import require_flash_attn, require_torch_accelerator, slow
+from transformers.testing_utils import Expectations, slow
 
 
 _TEST_PROMPTS = [
@@ -15,18 +15,27 @@
     "For safety reasons, the train is stopped in the middle of the",
 ]
 
-_EXPECTED_OUTPUTS = [
-    "a woman standing on the sidewalk, looking at him. He is immediately drawn to her and feels a strong attraction. He walks up to her and strikes up a conversation, and they quickly discover that they have a lot in common. They exchange numbers and",
-    "orange.\n\n## Step 1: Identify the key characteristics of the fruit\nThe fruit is described as being orange in color and round in shape.\n\n## Step 2: Determine the taste and nutritional value of the fruit\nThe fruit is described as sweet",
-    "This riddle is a classic example of a lateral thinking puzzle, which requires the test-taker to think creatively and consider multiple possibilities. The answer is not a straightforward one, and it requires some lateral thinking to arrive at the correct solution.",
-    "get in touch with us. We will respond to your message as soon as possible.\n\n[Your Name]\n[Your Email]\n[Your Phone Number]\n[Your Message]\n\nWe are looking forward to hearing from you!\n\n[Insert Contact Information]\n\nNote:",
-    "track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers does the train travel in 30 minutes?\n## Step 1: Convert the speed from km/h to km/min",
-]
+_EXPECTED_OUTPUTS = Expectations(
+    {
+        ("cpu", None): [
+            "orange.\n\n## Step 1: Identify the key characteristics of the fruit\nThe fruit is described as being orange in color and round in shape.\n\n## Step 2: Determine the taste and nutritional value of the fruit\nThe fruit is described as sweet",
+            "get started with our services.\nWe will be in touch with you shortly to discuss your project and provide a quote.\n\n**Project Details**\n\n* Project Name: _____________________________________________________\n* Project Description: __________________________________________________\n* Project Type (check all that apply",
+            "track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers does the train travel in 30 minutes?\n## Step 1: Convert the speed from km/h to km/min",
+            "This riddle is a classic example of a lateral thinking puzzle, which requires the test-taker to think creatively and consider multiple possibilities. The answer is not a straightforward one, and it requires some lateral thinking to arrive at the correct solution.",
+            "a woman standing on the sidewalk, looking at him. He is immediately drawn to her and feels a strong attraction. He walks up to her and strikes up a conversation.",
+        ],
+        (None, None): [
+            "a woman standing on the sidewalk, looking at him. He is immediately drawn to her and feels a strong attraction. He walks up to her and strikes up a conversation, and they quickly discover that they have a lot in common. They exchange numbers and",
+            "orange.\n\n## Step 1: Identify the key characteristics of the fruit\nThe fruit is described as being orange in color and round in shape.\n\n## Step 2: Determine the taste and nutritional value of the fruit\nThe fruit is described as sweet",
+            "This riddle is a classic example of a lateral thinking puzzle, which requires the test-taker to think creatively and consider multiple possibilities. The answer is not a straightforward one, and it requires some lateral thinking to arrive at the correct solution.",
+            "get in touch with us. We will respond to your message as soon as possible.\n\n[Your Name]\n[Your Email]\n[Your Phone Number]\n[Your Message]\n\nWe are looking forward to hearing from you!\n\n[Insert Contact Information]\n\nNote:",
+            "track. The train is stopped for 30 minutes. The train is moving at a speed of 60 km/h. How many kilometers does the train travel in 30 minutes?\n## Step 1: Convert the speed from km/h to km/min",
+        ],
+    }
+)
 
 
 @slow
-@require_flash_attn
-@require_torch_accelerator
 class TestBatchGeneration(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
@@ -77,11 +86,13 @@ def test_generate_batch_consistency(self, attn_impl, num_blocks, block_size, max
             f"\n[{attn_impl}] Batch took {end - start:.2f}s with config: blocks={num_blocks}, block_size={block_size}, max_batch_tokens={max_batch_tokens}"
         )
 
+        expected_outputs = _EXPECTED_OUTPUTS.get_expectation()
+
         for i, req_id in enumerate(batch_outputs):
             generated = self.tokenizer.decode(
                 batch_outputs[req_id].generated_tokens, skip_special_tokens=False
             ).strip()
-            expected = _EXPECTED_OUTPUTS[i].strip()
+            expected = expected_outputs[i].strip()
             self.assertTrue(
                 generated.startswith(expected),
                 msg=f"[{attn_impl}] Mismatch in request {i}:\nExpected start: {expected}\nGot: {generated}",