Repository: huggingface/transformers
Pull request: [CB] Easy optimizations for continuous batching

## Summary
This PR adds a few optimizations for continuous batching:
- removes non-needed `torch.cuda.synchronize` 
- add sorting of the inputs to maximize prefix caching hits
- sampling is done on the GPU
- removed an extranuous axis from the output_ids

## Performance 
| Attention         | Version       | Generated tokens | Duration (s) | Throughput (tok/s) |
|------------------------|----------------|--------------------------|------------------|--------------------------|
| Flash attention 3 | This PR         | 112599 | 16.73       | 6729.27           |
| Flash attention 3 | Main branch  | 111823 | 25.67       |  4355.68          |
| Flash attention 2 |  This PR        | 112822 | 24.61       | 4584.74          |
| Flash attention 2 | Main branch  | 112126 | 33.12       | 3385.46          |
| SDPA                  | This PR         | 113254  | 82.49       | 1373.00         |
| SDPA                  | Main branch  | 113725  | 170.84    | 665.67            |
* the performances on `main` were lowered by #42699 when compared to #42516 but this PR fixes this.

## Tests
No new failures for the CB tests. The only two failing tests are the same as in #42699 because of compile not working with `gemma2` which seems out of scope and acceptable.

Changed files:
filename: examples/pytorch/continuous_batching.py
patch: @@ -182,13 +182,16 @@ def batch_generate(
 
     # Benchmark parameters
     parser.add_argument("--samples", type=int, default=500, help="Number of samples to generate")
+    parser.add_argument(
+        "--input-length", type=int, default=None, help="Length of input sequences. Leave to None to mimic real eval."
+    )
     parser.add_argument("--max-new-tokens", type=int, default=512, help="Maximum number of new tokens to generate")
+    parser.add_argument("--force-max-length", action="store_true", help="Force generation to stop at max length")
 
     parser.add_argument("--add-prefix", action="store_true", help="Add a prefix to the samples")
     parser.add_argument("--compare", action="store_true", help="Compare CB generation with classic generate")
     parser.add_argument("--profile", type=str, default=None)
     parser.add_argument("--metrics", action="store_true")
-    parser.add_argument("--force-max-length", action="store_true", help="Force generation to stop at max length")
 
     # Display parameters
     parser.add_argument("--displayed", type=int, default=0, help="Number of samples to display")
@@ -251,6 +254,12 @@ def batch_generate(
     else:
         possible_prefixes = [None]
 
+    tokenizer_kwargs = {"add_generation_prompt": True}
+    if args.input_length is not None:
+        tokenizer_kwargs["max_length"] = args.input_length
+        tokenizer_kwargs["truncation"] = True
+        tokenizer_kwargs["padding"] = True
+
     batched_inputs = []
     for item, prefix in zip(dataset, cycle(possible_prefixes)):
         messages = []
@@ -261,7 +270,7 @@ def batch_generate(
             else:
                 question = prefix + "\n\n" + question
         messages.append({"role": "user", "content": question})
-        inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+        inputs = tokenizer.apply_chat_template(messages, **tokenizer_kwargs)
         inputs = inputs if isinstance(inputs, list) else inputs["input_ids"]
         batched_inputs.append(inputs)
 
@@ -283,6 +292,7 @@ def batch_generate(
         generation_cfg.compile_config = CompileConfig(
             fullgraph=True,
             mode="max-autotune-no-cudagraphs",
+            dynamic=True,  # FIXME: if we warmup all graphs, this is not needed anymore
         )
 
     # If we need to compare, we need to generate the reference outputs

filename: src/transformers/generation/continuous_batching/continuous_api.py
patch: @@ -259,7 +259,7 @@ def setup_static_tensors(self, num_groups: int) -> None:
         self.cumulative_seqlens_q = torch.empty((self.max_batch_tokens + 1,), **self.tensor_metadata)
         self.max_seqlen_q = 0
         self.logits_indices = torch.empty((self.max_batch_tokens,), **self.tensor_metadata)
-        self.output_ids = torch.empty((1, self.max_batch_tokens), **self.tensor_metadata)
+        self.output_ids = torch.empty((self.max_batch_tokens,), **self.tensor_metadata)
 
         # For some kwargs, we have a dict of tensors with as many items as there are attention types
         layer_types = getattr(self.config, "layer_types", None)
@@ -311,7 +311,7 @@ def reset_static_tensors(self, full_reset: bool = False) -> None:
         self.cumulative_seqlens_q[: b_size + 1].zero_()
         self.max_seqlen_q = 0
         self.logits_indices[:q_len].fill_(-1)
-        self.output_ids[:, :q_len].fill_(-1)
+        self.output_ids[:q_len].fill_(-1)
 
         # Reset the attributes that are either tensors or dict of tensors
         for layer_type in self.cumulative_seqlens_k:
@@ -447,7 +447,7 @@ def prepare_next_batch(self) -> bool:
         self.metrics.record_batch_metrics(self.requests_in_batch)
 
         # Reset the static tensors used for storage
-        self.reset_static_tensors()  # TODO: this might be unnecessary
+        self.reset_static_tensors()  # FIXME: why does this make the generation faster?
 
         # Prepare accumulators
         self.actual_query_length = 0
@@ -557,13 +557,10 @@ def _build_tensors(
             self.actual_index_sizes[i] = (len(group_read_indices), len(group_write_indices))
 
     @traced
-    def _sync(self) -> list[int]:
-        if self.output_ids is not None:
-            try:
-                return self.output_ids.tolist()[0]
-            except Exception:
-                return [0, 1]
-        return [0, 0]
+    def _get_new_tokens(self, num_new_tokens: int) -> list[int]:
+        indices = self.logits_indices[:num_new_tokens]
+        new_tokens = self.output_ids[indices]
+        return new_tokens.tolist()
 
     @traced
     def _maybe_send_output(self, state: RequestState) -> None:
@@ -574,13 +571,13 @@ def _maybe_send_output(self, state: RequestState) -> None:
     @traced
     def update_batch(self) -> None:
         """Update request states based on generated tokens."""
-        out_tokens = self._sync()
+        new_tokens = self._get_new_tokens(len(self.requests_in_batch))
         for i, state in enumerate(self.requests_in_batch):
             # If the request has no remaining prompt ids, it means prefill has already ended or just finished
             if len(state.remaining_prefill_tokens) == 0:
                 self.metrics.record_ttft_metric(state.created_time, state.request_id)
                 state.status = RequestStatus.DECODING
-                token = out_tokens[self.logits_indices[i]]
+                token = new_tokens[i]
                 state.tokens_to_process = [token]
                 # Update the request and stop if it is complete
                 is_finished = state.update_and_check_completion(token)
@@ -727,12 +724,11 @@ def _sample(self, probs: torch.Tensor, do_sample: bool) -> None:
             probs = nn.functional.softmax(probs, dim=-1)
             # probs[0] has shape [seq_len, vocab_size], multinomial returns [seq_len, 1]
             next_tokens = torch.multinomial(probs[0], num_samples=1).squeeze(-1)  # Now [seq_len]
-            # Add batch dimension back to match argmax output
-            next_tokens = next_tokens.unsqueeze(0)  # Now [1, seq_len]
         else:
-            next_tokens = torch.argmax(probs, dim=-1)  # Already [1, seq_len]
-        tokens = next_tokens.size(1)  # Get seq_len dimension
-        self.output_ids[:, :tokens].copy_(next_tokens)
+            next_tokens = torch.argmax(probs, dim=-1)  # shape is [1, seq_len]
+            next_tokens = next_tokens.squeeze(0)  # shape is [seq_len]
+        tokens = next_tokens.size(0)  # Get seq_len dimension
+        self.output_ids[:tokens].copy_(next_tokens)
 
 
 # Manager Class (User Interface)
@@ -950,6 +946,10 @@ def add_requests(
         streaming: bool = False,
         record_timestamps: bool = False,
     ) -> None:
+        # If there is prefix sharing, we sort the inputs to maximize cache hits
+        if self._allow_prefix_sharing:
+            inputs = sorted(inputs, reverse=True)
+        # Add requests in order
         for input_ids in inputs:
             self.add_request(
                 input_ids, max_new_tokens=max_new_tokens, streaming=streaming, record_timestamps=record_timestamps
@@ -1080,10 +1080,6 @@ def _inner_generation_loop(self, batch_processor: ContinuousBatchProcessor) -> N
             )
 
         self._generation_step()
-
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()  # FIXME: why is this needed?
-        # Processor updates the batch after generation step is truly over
         batch_processor.update_batch()
 
     @traced

filename: tests/generation/test_continuous_batching.py
patch: @@ -251,7 +251,17 @@ def _test_continuous_batching_parity(
         generate_outputs = model.generate(**inputs.to(torch_device), generation_config=model.generation_config)
 
         for i, user_message in enumerate(user_messages):
-            continuous_batching_output = continuous_batching_outputs[f"req_{i}"].generated_tokens
+            # Find the corresponding request in the continuous batching outputs
+            input_tokens = inputs.input_ids[i][inputs.attention_mask[i] == 1].tolist()
+            key_to_pop = None
+            for key, state in continuous_batching_outputs.items():
+                if state.prompt_ids == input_tokens:
+                    key_to_pop = key
+                    break
+            if key_to_pop is None:
+                self.fail(f"Request {i} not found in continuous batching outputs")
+            continuous_batching_output = continuous_batching_outputs.pop(key_to_pop).generated_tokens
+
             generate_output = generate_outputs[i][num_input_tokens:].tolist()
             while generate_output[-1] == model.generation_config.pad_token_id:
                 generate_output.pop()