Repository: huggingface/transformers
Pull request: Streaming should be handled at the request-level rather than at the istance level


Changed files:
filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -149,7 +149,6 @@ def __init__(
         model_device: torch.device,
         model_dtype: torch.dtype,
         scheduler: Scheduler,
-        streaming: bool = False,
         manual_eviction: bool = False,
         slice_inputs: bool = True,  # TODO: There should be an heuristic to decide on slicing, compile, cuda graphs...
     ) -> None:
@@ -165,7 +164,6 @@ def __init__(
             model_device: Device for model inputs/outputs
             model_dtype: Data type for model inputs/outputs
             scheduler: The [`Scheduler`] to use
-            streaming: Whether to stream tokens as they're generated
             manual_eviction: Whether to manually evict blocks from the cache
             slice_inputs: Whether to slice the inputs to the model
         """
@@ -178,7 +176,6 @@ def __init__(
         self.model_device = model_device
         self.model_dtype = model_dtype
         self.scheduler = scheduler
-        self.streaming = streaming
         self.manual_eviction = manual_eviction
         self.slice_inputs = slice_inputs
 
@@ -519,7 +516,7 @@ def _sync(self):
     @traced
     def _maybe_send_output(self, state: RequestState, token: int):
         """Send output to the queue based on streaming mode and request state."""
-        if self.streaming:
+        if state.streaming:
             self.output_queue.put(state.to_generation_output())
         elif state.status == RequestStatus.FINISHED:
             self.output_queue.put(state.to_generation_output())
@@ -596,7 +593,6 @@ def __init__(
         generation_config: GenerationConfig,
         manual_eviction: bool = False,
         max_queue_size=0,
-        streaming: bool = True,
         slice_inputs: bool = True,
     ):
         """
@@ -606,7 +602,6 @@ def __init__(
             model: The language model for generation
             generation_config: Configuration for generation parameters
             max_queue_size: Maximum size of the request queue (0 = unlimited)
-            streaming: Whether to stream tokens as they are generated
         """
         if "paged|" not in model.config._attn_implementation:
             attn_implementation = f"paged|{model.config._attn_implementation}"
@@ -625,7 +620,6 @@ def __init__(
         self.input_queue = queue.Queue(maxsize=max_queue_size)
         self.output_queue = queue.Queue()
         self.stop_event = threading.Event()
-        self.streaming = streaming
         self.log_prob_generation = getattr(generation_config, "log_prob_generation", False)
         self._generation_thread = None
         self._request_counter = 0
@@ -690,7 +684,11 @@ def join(self, timeout: Optional[float] = None):
                 self._generation_thread = None
 
     def add_request(
-        self, input_ids: list[int], request_id: Optional[str] = None, max_new_tokens: Optional[int] = None
+        self,
+        input_ids: list[int],
+        request_id: Optional[str] = None,
+        max_new_tokens: Optional[int] = None,
+        streaming: bool = False,
     ) -> str:
         """Add a new generation request to the queue.
 
@@ -716,6 +714,7 @@ def add_request(
             full_prompt_ids=list(input_ids),
             max_new_tokens=max_new_tokens,
             eos_token_id=self.generation_config.eos_token_id,
+            streaming=streaming,
         )
 
         # Use block=True with timeout to handle backpressure if queue is full
@@ -872,7 +871,6 @@ def _run_generation_loop(self):
                 self.model.device,
                 self.model.dtype,
                 scheduler(paged_attention_cache, self.manual_eviction),
-                self.streaming,
                 self.manual_eviction,
                 slice_inputs=self.slice_inputs,
             )
@@ -956,15 +954,13 @@ def init_continuous_batching(
         generation_config: Optional[GenerationConfig] = None,
         manual_eviction: bool = False,
         max_queue_size: int = 0,
-        streaming: bool = False,
         slice_inputs: bool = True,
     ) -> ContinuousBatchingManager:
         """Initialize a manager for continuous batching inference.
 
         Args:
             generation_config: Custom generation configuration
             max_queue_size: Maximum size of the input request queue
-            streaming: Whether to stream tokens as they are generated
 
         Returns:
             `ContinuousBatchingManager`: The manager instance to add requests and retrieve results.
@@ -986,7 +982,6 @@ def init_continuous_batching(
             generation_config=gen_config,
             manual_eviction=manual_eviction,
             max_queue_size=max_queue_size,
-            streaming=streaming,
             slice_inputs=slice_inputs,
         )
 

filename: src/transformers/generation/continuous_batching/requests.py
patch: @@ -102,6 +102,7 @@ class RequestState:
                                 SPLIT_PENDING_REMAINDER, DECODING, FINISHED, FAILED
         max_new_tokens (int): The maximum number of new tokens to generate.
         eos_token_id (int): The ID of the end-of-sequence token.
+        streaming (bool): Whether to stream tokens as they're generated
         created_time (float): The time the request was created.
         error (Optional[str]): Any error message associated with the request. When None, has had no error yet.
     """
@@ -117,6 +118,7 @@ class RequestState:
     _status: RequestStatus = RequestStatus.PENDING  # Status of the request, hidden behind a property
     max_new_tokens: int = 20  # Maximum number of new tokens to generate
     eos_token_id: int = -1  # ID of the end-of-sequence token
+    streaming: bool = False  # Whether to stream tokens as they're generated
     created_time: float = field(default_factory=time.time)  # Time the request was created
     error: Optional[str] = None  # Error message if the request failed
     lifespan: tuple[float, float] = (-1, -1)  # (time request was no longer pending, time request finished)

filename: tests/generation/test_continuous_batching.py
patch: @@ -18,7 +18,7 @@
 import torch
 from parameterized import parameterized
 
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList
 from transformers.generation.continuous_batching.cache import group_layers_by_attn_type
 from transformers.generation.continuous_batching.continuous_api import build_attention_mask
 from transformers.testing_utils import Expectations, require_kernels, require_torch_gpu, slow
@@ -337,6 +337,102 @@ def test_attn_implementation(self) -> None:
         manager = model.init_continuous_batching()
         assert "paged|eager" == manager.model.config._attn_implementation
 
+    @require_torch_gpu
+    def test_streaming_request(self) -> None:
+        model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+        max_new_tokens = 3
+
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+
+        manager = model.init_continuous_batching()
+        manager.logit_processor = LogitsProcessorList()
+        manager.start()
+
+        messages = [{"content": "What is the Transformers library known for?", "role": "user"}]
+
+        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(
+            model.device
+        )[0]
+
+        request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=True)
+
+        # In streaming mode, the total number of generated tokens is incremented by 1 on each iteration
+        chunk_1 = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk_1.generated_tokens), 1)
+
+        chunk_2 = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk_2.generated_tokens), 2)
+
+        chunk_3 = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk_3.generated_tokens), 3)
+
+        manager.stop(block=True)
+
+    @require_torch_gpu
+    def test_non_streaming_request(self) -> None:
+        model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+        max_new_tokens = 3
+
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+
+        manager = model.init_continuous_batching()
+        manager.logit_processor = LogitsProcessorList()
+        manager.start()
+
+        messages = [{"content": "What is the Transformers library known for?", "role": "user"}]
+
+        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(
+            model.device
+        )[0]
+
+        request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=False)
+
+        chunk = next(manager.request_id_iter(request_id))
+
+        # In non-streaming mode, the total number of generated tokens is equal to the max new tokens
+        self.assertEqual(len(chunk.generated_tokens), max_new_tokens)
+
+        manager.stop(block=True)
+
+    @require_torch_gpu
+    def test_streaming_and_non_streaming_requests_can_alternate(self) -> None:
+        model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+        max_new_tokens = 3
+
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
+
+        manager = model.init_continuous_batching()
+        manager.logit_processor = LogitsProcessorList()
+        manager.start()
+
+        messages = [{"content": "What is the Transformers library known for?", "role": "user"}]
+
+        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(
+            model.device
+        )[0]
+
+        # Non-streaming request
+        request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=False)
+        chunk = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk.generated_tokens), max_new_tokens)
+
+        # Streaming request works afterward
+        request_id = manager.add_request(inputs, max_new_tokens=max_new_tokens, streaming=True)
+
+        chunk_1 = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk_1.generated_tokens), 1)
+
+        chunk_2 = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk_2.generated_tokens), 2)
+
+        chunk_3 = next(manager.request_id_iter(request_id))
+        self.assertEqual(len(chunk_3.generated_tokens), 3)
+
+        manager.stop(block=True)
+
 
 # FIXME: the gemma test seem broken, there is a message about cuda graphs and the sdpa and flash expecteations are
 # inverted on CUDA. On AMD they do fine.