Repository: huggingface/transformers
Pull request: Switch to CB if cache_implementation == paged

This PR fixes one of the issues raised in #41501

This PR adds `paged` as a static cache implementation and switches from `generate` to `generate_batch` if it's the chosen implementation. It also raises errors if an argument for `generate` is passed to `generate_batch` but is not supported, and warn the user that `generate_batch` is used. 
The interface for the user stays the same, as the sample script below shows:
```
import os
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig


if __name__ == '__main__':
    model_path = "Qwen/Qwen3-0.6B"

    tokenizer = AutoTokenizer.from_pretrained(
        model_path,
        trust_remote_code=True,
        add_bos_token=False,
        add_eos_token=False,
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map="cuda",
    ).eval()

    prompt="How are you today?"
    chat_prompt=tokenizer.apply_chat_template(
        [{"role" : "user", "content" : prompt}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    generation_config=GenerationConfig(
        do_sample=True,
        temperature=1.0,
        top_p=1.0,
        # num_return_sequences=8,
        max_new_tokens=100,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        use_cache=True,
        return_dict_in_generate=True,
        output_logits=True,
    )

    chat_prompt_ids = tokenizer(chat_prompt, return_tensors="pt")["input_ids"].to(model.device)

    print("Classic generate")
    output = model.generate(input_ids=chat_prompt_ids, generation_config=generation_config)[0]
    output_text = tokenizer.batch_decode(output[:, chat_prompt_ids.shape[1]:], skip_special_tokens=True)
    print(output_text)

    print("\nWith paged cache")
    output = model.generate(input_ids=chat_prompt_ids, generation_config=generation_config, cache_implementation="paged")[0]
    output_text = tokenizer.batch_decode(output[:, chat_prompt_ids.shape[1]:], skip_special_tokens=True)
    print(output_text)

```
which outputs
```
Classic generate
`generation_config` default values have been modified to match model-specific defaults: {'temperature': 0.6, 'top_k': 20, 'top_p': 0.95, 'bos_token_id': 151643}. If this is not desired, please set these values explicitly.
["Hello! I'm here to help you with anything you need. How can I assist you today?"]

With paged cache
Detected cache_implementation=paged: switching to continuous batching. You should consider using generate_batch directly instead.
No behavior specified for use_cuda_graph, defaulting to self.use_cuda_graph = False because self.model.config._attn_implementation = 'paged|sdpa'. If you want to save memory, turn off cuda graphs, but they can improve performances.
Solving 1 requests: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  2.16request/s]
["Hello! I'm here to help you with anything you need. How can I assist you today? 😊"]
```

Changed files:
filename: src/transformers/generation/configuration_utils.py
patch: @@ -556,10 +556,13 @@ def validate(self, strict=False):
                 "`model.generation_config.pad_token_id=PAD_TOKEN_ID` to avoid errors in generation"
             )
         # 1.2. Cache attributes
-        if self.cache_implementation is not None and self.cache_implementation not in ALL_CACHE_IMPLEMENTATIONS:
+        # "paged" re-routes to continuous batching and so it is a valid cache implementation. But we do not want to test
+        # it with the `generate` as the other would be, so we we cannot add it to ALL_CACHE_IMPLEMENTATIONS
+        valid_cache_implementations = ALL_CACHE_IMPLEMENTATIONS + ("paged",)
+        if self.cache_implementation is not None and self.cache_implementation not in valid_cache_implementations:
             raise ValueError(
                 f"Invalid `cache_implementation` ({self.cache_implementation}). Choose one of: "
-                f"{ALL_CACHE_IMPLEMENTATIONS}"
+                f"{valid_cache_implementations}"
             )
         # 1.3. Performance attributes
         if self.compile_config is not None and not isinstance(self.compile_config, CompileConfig):

filename: src/transformers/generation/utils.py
patch: @@ -2370,7 +2370,7 @@ def generate(
                     - [`~generation.GenerateEncoderDecoderOutput`],
                     - [`~generation.GenerateBeamEncoderDecoderOutput`]
         """
-        # 0. If requested, load an arbitrary generation recipe from the Hub and run it instead
+        # 0.a. If requested, load an arbitrary generation recipe from the Hub and run it instead
         trust_remote_code = kwargs.pop("trust_remote_code", None)
 
         if custom_generate is not None and isinstance(custom_generate, str):
@@ -2392,6 +2392,75 @@ def generate(
             )
             return custom_generate_function(model=self, **generate_arguments)
 
+        # 0.b. If requested, switched to continuous batching generation
+        if kwargs.get("cache_implementation") == "paged":
+            logger.warning(
+                "Detected cache_implementation=paged: switching to continuous batching. You should consider using "
+                "generate_batch directly instead."
+            )
+
+            # generate_batch expects a list of lists of ints, so we create it from the inputs or input_ids
+            inputs = inputs if inputs is not None else kwargs.get("input_ids")
+            if inputs is None:
+                raise ValueError("inputs or input_ids must be provided for CB generation.")
+
+            if inputs.dim() == 1:
+                inputs = inputs.unsqueeze(0).tolist()
+            elif inputs.dim() == 2:
+                inputs = inputs.tolist()
+            else:
+                raise ValueError(f"inputs must be a 1D or 2D tensor, got {inputs.dim() = }")
+
+            # some arguments are not supported for continuous batching
+            if stopping_criteria is not None:
+                raise NotImplementedError(
+                    f"stopping_criteria is not supported for continuous batching. Got {stopping_criteria = }"
+                )
+            if prefix_allowed_tokens_fn is not None:
+                raise NotImplementedError(
+                    f"prefix_allowed_tokens_fn is not supported for continuous batching. Got {prefix_allowed_tokens_fn = }"
+                )
+            if assistant_model is not None:
+                raise NotImplementedError(
+                    f"assistant_model is not supported for continuous batching. Got {assistant_model = }"
+                )
+            if streamer is not None:  # TODO: actualy this could be supported
+                raise NotImplementedError(f"streaming is not supported for continuous batching. Got {streamer = }")
+            if negative_prompt_ids is not None:
+                raise NotImplementedError(
+                    f"negative_prompt_ids is not supported for continuous batching. Got {negative_prompt_ids = }"
+                )
+            if negative_prompt_attention_mask is not None:
+                raise NotImplementedError(
+                    f"negative_prompt_attention_mask is not supported for continuous batching. Got {negative_prompt_attention_mask = }"
+                )
+
+            # others are ignored
+            if synced_gpus is not None:
+                logger.warning(f"synced_gpus is not ignored for continuous batching. Got {synced_gpus = }")
+            num_return_sequences = kwargs.get("num_return_sequences", 1)
+            num_beams = kwargs.get("num_beams", 1)
+            if num_return_sequences > 1 or num_beams > 1:  # FIXME: remove this once CB supports it (which is planned)
+                logger.warning(
+                    f"num_return_sequences and num_beams are not supported for continuous batching yet. "
+                    f"Got {num_return_sequences = } and {num_beams = }. "
+                )
+
+            # switch to CB
+            outputs = self.generate_batch(
+                inputs=inputs,
+                generation_config=self._prepare_generation_config(generation_config, use_model_defaults, **kwargs)[0],
+                **kwargs,
+            )
+            sequences = [
+                outputs[f"req_{i}"].prompt_ids + outputs[f"req_{i}"].generated_tokens for i in range(len(outputs))
+            ]
+
+            # To use the same indexing (outputs[0]) as the regular generate method, we unsqueeze the tensor
+            sequences_as_tensor = torch.tensor(sequences, dtype=torch.long, device=self.device)
+            sequences_as_tensor = sequences_as_tensor.unsqueeze(0)
+            return sequences_as_tensor
+
         # 1. Handle kwargs, `generation_config`, validate them and obtain generation mode
         generation_mode_kwargs = self._extract_generation_mode_kwargs(
             custom_generate,