transformers 5.0.0__py3-none-any.whl → 5.0.0rc0__py3-none-any.whl

transformers/generation/continuous_batching/continuous_api.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2024 The HuggingFace Inc. team.
 # Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
 #
@@ -28,12 +29,12 @@ from tqdm import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
 
 from ...configuration_utils import PretrainedConfig
-from ...generation.configuration_utils import CompileConfig, GenerationConfig
+from ...generation.configuration_utils import GenerationConfig
 from ...generation.logits_process import LogitsProcessor
 from ...utils.logging import logging
 from ...utils.metrics import ContinuousBatchProcessorMetrics, attach_tracer, traced
 from .cache import PagedAttentionCache
-from .requests import GenerationOutput, RequestState, RequestStatus, logger
+from .requests import GenerationOutput, RequestState, RequestStatus, get_device_and_memory_breakdown, logger
 from .scheduler import SCHEDULER_MAPPING, FIFOScheduler, Scheduler
 
 
@@ -44,20 +45,17 @@ generation goes on, there are two dimensions that change:
 - the number of keys/values tokens (KV), which grows as the cache does
 
 To solve this, we slice along those dimensions to fixed lengths. The size of the slices is controlled by the variables
-num_x_padding_intervals: NUM_X_PADDING_INTERVALS means that we create at most NUM_X_PADDING_INTERVALS graphs for the X
-dimension. So if the maximum number of queries tokens is 1000, and NUM_Q_PADDING_INTERVALS is 4, we will slice the
-number of queries token by intervals of 1000 / 4 = 250 tokens, ie. to 250, 500, 750 or 1000 queries tokens.
+below: NUM_X_CUDA_GRAPHS means that we create at most NUM_X_CUDA_GRAPHS graphs for the X dimension. So if the maximum
+number of queries tokens is 1000, and NUM_Q_CUDA_GRAPHS is 4, we will slice the number of queries token by intervals of
+1000 / 4 = 250 tokens, ie. to 250, 500, 750 or 1000 queries tokens.
 
 Smaller slices means more granularity and thus less padding. But since each graph takes up space on the GPU and time to
 create, we don't want to many graphs. And since the size of the KV dimension is the number of queries tokens plus the
 number of tokens cached, dimension of KV is usually much larger than the dimension of Q. So we have more granularity
 for the KV dimension than the query dimension.
-
-This variable used to be called NUM_X_CUDA_GRAPHS, but we renamed it to NUM_X_PADDING_INTERVALS because it is used for
-padding in the case of cuda graphs AND torch.compile.
 """
-NUM_Q_PADDING_INTERVALS = 4
-NUM_KV_PADDING_INTERVALS = 8
+NUM_Q_CUDA_GRAPHS = 4
+NUM_KV_CUDA_GRAPHS = 8
 
 
 def pad_by_intervals(size: int, max_value: int, nb_intervals: int) -> int:
@@ -65,7 +63,7 @@ def pad_by_intervals(size: int, max_value: int, nb_intervals: int) -> int:
     interval_size = max_value // nb_intervals
     if interval_size == 0:
         return max_value
-    padded = ceil(size / interval_size) * interval_size if size > 0 else interval_size
+    padded = ceil(size / interval_size) * interval_size
     return min(padded, max_value)
 
 
@@ -190,8 +188,6 @@ class ContinuousBatchProcessor:
         scheduler: Scheduler,
         manual_eviction: bool,
         use_cuda_graph: bool,
-        q_padding_intervals: int,
-        kv_padding_intervals: int,
     ) -> None:
         """Initialize the continuous batch processor.
 
@@ -225,14 +221,7 @@ class ContinuousBatchProcessor:
         # Accumulator for batch scheduling
         self.requests_in_batch: list[RequestState] = []
         # Cuda graphs for the generation step
-        self.q_padding_intervals = q_padding_intervals
-        self.kv_padding_intervals = kv_padding_intervals
         self._graphs: dict[tuple[int, int], torch.cuda.CUDAGraph] | None = {} if use_cuda_graph else None
-        # Compile-related arguments
-        self.compile_config: CompileConfig | None = getattr(generation_config, "compile_config", None)
-        self._forward_process_and_sample_is_compiled = False
-
-        self._pad_inputs = use_cuda_graph or (self.compile_config is not None and not self.compile_config.dynamic)
 
         # Set up metrics collector
         self.max_batch_tokens = cache.max_batch_tokens
@@ -249,7 +238,7 @@ class ContinuousBatchProcessor:
     def setup_static_tensors(self, num_groups: int) -> None:
         """Setup the static tensors that are used for storage during the generation step. No other tensor will be
         allowed for the inputs or the outputs of the generation step."""
-        self.num_pages = self.cache.num_blocks * self.cache.block_size
+        num_pages = self.cache.num_blocks * self.cache.block_size
         self.tensor_metadata = {"dtype": torch.int32, "device": self.model_device}
 
         # Some tensors always have the same shape regardless of the model
@@ -258,7 +247,7 @@ class ContinuousBatchProcessor:
         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((self.max_batch_tokens,), **self.tensor_metadata)
+        self.output_ids = torch.empty((1, 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)
@@ -274,7 +263,7 @@ class ContinuousBatchProcessor:
 
         if attn_mask_is_needed(self.config):
             attn_mask_kwargs = {
-                "size": (1, 1, self.max_batch_tokens, self.num_pages + self.max_batch_tokens),
+                "size": (1, 1, self.max_batch_tokens, num_pages + self.max_batch_tokens),
                 "dtype": self.model_dtype,
                 "device": self.model_device,
             }
@@ -287,7 +276,7 @@ class ContinuousBatchProcessor:
             torch.empty((self.max_batch_tokens,), **self.tensor_metadata) for _ in range(num_groups)
         ]
         self.read_index_storage = [
-            torch.empty((self.num_pages + self.max_batch_tokens), **self.tensor_metadata) for _ in range(num_groups)
+            torch.empty((num_pages + self.max_batch_tokens), **self.tensor_metadata) for _ in range(num_groups)
         ]
         # For read index, the +T is because there are -1 for seqlen_q when model uses a sliding window
 
@@ -310,7 +299,7 @@ class ContinuousBatchProcessor:
         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:
@@ -427,33 +416,6 @@ class ContinuousBatchProcessor:
         self.metrics.record_request_completion(state.created_time, state.request_id)
         self.output_queue.put(state.to_generation_output())
 
-    # TODO: there should be a way to choose the offloading policy: biggest request, oldest request, etc.
-    # Including a policy to not allow offloading and crashing the generation
-    def soft_reset_one_request(self) -> None:
-        """Soft resets one active request by removing it from active requests and re-adding it to the waiting queue.
-
-        The generated tokens are kept as part of the new request's initial prompt. When `block_new_requests` is False,
-        the oldest request is offloaded; when True, the newest request is offloaded. This method also sets
-        `block_new_requests` to True to prevent infinite loops of offloading and re-scheduling requests.
-        """
-        # The offloaded request is the newest (resp. oldest) if block_new_requests is True (resp. False)
-        if self.scheduler.block_new_requests:
-            request_id, state = self.scheduler.active_requests.popitem()
-        else:
-            request_id, state = next(iter(self.scheduler.active_requests.items()))
-        logger.info(
-            f"Soft resetting request {request_id} with {len(state.initial_tokens)} initial tokens and "
-            f"{len(state.generated_tokens)} generated tokens"
-        )
-        # Create a copy of the offloaded request keeping the generated tokens as addition to the initial prompt
-        new_state = state.create_equivalent_initial_request()
-        # Actual offloading of the request
-        self.scheduler.finish_request(request_id, evict_from_cache=True)
-        self.scheduler.add_waiting_request(new_state)
-        # This flag blocks any new requests from being scheduled until one request is finished. This ensures that we
-        # don't enter an offload / schedule loop
-        self.scheduler.block_new_requests = True
-
     @traced
     def prepare_next_batch(self) -> bool:
         """Prepare tensors and metadata for the next model forward pass. Returns True if there are requests to process,
@@ -467,22 +429,13 @@ class ContinuousBatchProcessor:
         self.metrics.record_queue_metrics(len(self.scheduler.active_requests), len(self.scheduler.waiting_requests))
 
         # Schedule the next batch of requests, stop if there are no requests in the batch
-        self.requests_in_batch = self.scheduler.schedule_batch(self.max_batch_tokens, self.num_pages)
-
-        # If requests_in_batch is None, it means we need to offload some requests if possible
-        if self.requests_in_batch is None:
-            if len(self.scheduler.active_requests) > 1:
-                self.soft_reset_one_request()
-            else:
-                raise RuntimeError("No requests can be scheduled and no request can be offloaded.")
-        # If it's an empty list, it means we have no requests to process
+        self.requests_in_batch = self.scheduler.schedule_batch(self.max_batch_tokens)
         if not self.requests_in_batch:
             return False
-        # Otherwise, we can continue with the non-empty batch
         self.metrics.record_batch_metrics(self.requests_in_batch)
 
         # Reset the static tensors used for storage
-        self.reset_static_tensors()  # FIXME: why does this make the generation faster?
+        self.reset_static_tensors()  # TODO: this might be unnecessary
 
         # Prepare accumulators
         self.actual_query_length = 0
@@ -592,10 +545,13 @@ class ContinuousBatchProcessor:
             self.actual_index_sizes[i] = (len(group_read_indices), len(group_write_indices))
 
     @traced
-    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()
+    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]
 
     @traced
     def _maybe_send_output(self, state: RequestState) -> None:
@@ -606,56 +562,31 @@ class ContinuousBatchProcessor:
     @traced
     def update_batch(self) -> None:
         """Update request states based on generated tokens."""
-        new_tokens = self._get_new_tokens(len(self.requests_in_batch))
-        current_logits_index = 0
-        for state in self.requests_in_batch:
+        out_tokens = self._sync()
+        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:
-                # If there are no generated tokens yet, it means prefill just ended
-                if state.generated_len() == 0:
-                    self.metrics.record_ttft_metric(state.created_time, state.request_id)
-                    state.status = RequestStatus.DECODING
-
-                token = new_tokens[current_logits_index]
+                self.metrics.record_ttft_metric(state.created_time, state.request_id)
+                state.status = RequestStatus.DECODING
+                token = out_tokens[self.logits_indices[i]]
                 state.tokens_to_process = [token]
-                current_logits_index += 1
-
                 # Update the request and stop if it is complete
                 is_finished = state.update_and_check_completion(token)
                 # We mark the completed blocks as such
-                self.cache.mark_shareable_blocks_as_complete(state)
+                self.cache.mark_blocks_as_complete(state)
                 if is_finished:
                     self.metrics.record_request_completion(state.created_time, state.request_id)
                     self.scheduler.finish_request(state.request_id, evict_from_cache=(not self.manual_eviction))
-                    self.scheduler.block_new_requests = False
                 self._maybe_send_output(state)
             #  Otherwise, the request is still prefilling, but the prefill has been split
             elif state.status == RequestStatus.PREFILLING_SPLIT:
-                self.cache.mark_shareable_blocks_as_complete(state)
+                self.cache.mark_blocks_as_complete(state)
                 state.status = RequestStatus.SPLIT_PENDING_REMAINDER
             else:
                 raise ValueError(f"Request {state.request_id} is in an unexpected state: {state.status}")
 
-        # If some requests need to be forked, we do it now
-        copy_source, copy_destination = [], []
-        while self.scheduler._requests_to_fork:
-            # Get the number of children and reset it so it's not forked again
-            state = self.scheduler._requests_to_fork.pop()
-            num_children = state.num_children
-            state.num_children = 0
-            # Create the new request and add them to the scheduler
-            new_request_ids = [f"{state.request_id}__child#{i}" for i in range(num_children)]
-            for new_request_id in new_request_ids:
-                self.scheduler.active_requests[new_request_id] = state.fork(new_request_id)
-            # Fork the cache
-            copy_src, copy_dst = self.cache.fork_request(state.request_id, new_request_ids)
-            copy_source.extend(copy_src)
-            copy_destination.extend(copy_dst)
-            # FIXME: if fork cant be done, create a new pending request without forking instead of crashing everything
-
-        # The copy induced by the fork is done in one go (if it's even needed)
-        if copy_source:
-            self.cache.copy_cache(copy_source, copy_destination)
+        if self.cache.get_num_free_blocks() == 0:
+            raise ValueError("No more free blocks")
 
     @traced
     def has_pending_requests(self) -> bool:
@@ -696,36 +627,28 @@ class ContinuousBatchProcessor:
     def _generation_step(self, model: nn.Module, logit_processor: LogitsProcessor, do_sample: bool) -> None:
         """Perform a single generation step."""
 
-        # If a compile config is specified, we compile the forward pass once in a wrapper
-        if self.compile_config is not None and not self._forward_process_and_sample_is_compiled:
-            self._forward_process_and_sample = torch.compile(
-                self._forward_process_and_sample,
-                fullgraph=self.compile_config.fullgraph,
-                mode=self.compile_config.mode,
-                dynamic=self.compile_config.dynamic,
-                backend=self.compile_config.backend,
-                options=self.compile_config.options,
-            )
-            self._forward_process_and_sample_is_compiled = True
-
-        # If inputs are static sized, we find the padded sizes of the queries and keys/values
-        if self._pad_inputs:
-            padded_q = pad_by_intervals(self.actual_query_length, self.max_batch_tokens, self.q_padding_intervals)
-            max_read_index_size = max(self.actual_index_sizes[i][0] for i in range(self.cache.num_groups))
-            # The space planned for query tokens will be added later, so we remove it from the space planned for KV
-            padded_read_index_size = pad_by_intervals(max_read_index_size, self.num_pages, self.kv_padding_intervals)
-        else:
-            padded_q, padded_read_index_size = 0, 0
-        # Retrieve the model kwargs with or without padding
-        batch_data = self.get_model_kwargs(padded_q, padded_read_index_size)
-
-        # If we are not using cuda graphs, we perform the generation step and return
+        # If cuda graphs are disabled, we just use the actual size
         if self._graphs is None:
+            batch_data = self.get_model_kwargs()
             self._forward_process_and_sample(model, batch_data, logit_processor, do_sample)
             return None
 
-        # If we have a graph that fits, we replay it
+        # Determine the padded size of the queries and keys/values
+        padded_q = pad_by_intervals(self.actual_query_length, self.max_batch_tokens, NUM_Q_CUDA_GRAPHS)
+
+        max_read_index_size = max(self.actual_index_sizes[i][0] for i in range(self.cache.num_groups))
+        padded_read_index_size = pad_by_intervals(
+            max_read_index_size - self.max_batch_tokens,
+            self.cache.num_blocks * self.cache.block_size,
+            NUM_KV_CUDA_GRAPHS,
+        )
+
+        # Get the batch data and the associated graph
+        batch_data = self.get_model_kwargs(padded_q, padded_read_index_size)
+
         graph = self._graphs.get((padded_q, padded_read_index_size))
+
+        # If we have a graph that fits, we replay it
         if graph is not None:
             graph.replay()
             return None
@@ -750,6 +673,7 @@ class ContinuousBatchProcessor:
     ) -> None:
         """This function performs the forward pass, logits processing, and sampling; which are broken down into smaller
         function to be easier to trace with OpenTelemetry."""
+        # with torch.no_grad():
         logits = self._model_forward(model, batch_data)
         # if self.log_prob_generation:    batch_processor.output_probs.copy_(logits)  # TODO
         probs = self._process_logit(batch_data, logits, logit_processor)
@@ -767,7 +691,6 @@ class ContinuousBatchProcessor:
         # Handle shape compatibility: logit processors expect 2D tensors [batch_size, vocab_size]
         # but continuous batching always produces 3D tensors [batch_size, seq_len, vocab_size]
         batch_size, seq_len, vocab_size = logits.shape
-        # NOTE: to be an exact match with generate, we should also convert logits2d to float32 here, but it's not needed in practice
         logits_2d = logits.view(batch_size * seq_len, vocab_size)
         input_ids_2d = batch_data["input_ids"].view(batch_size * seq_len)
         # Process with 2D tensors
@@ -781,11 +704,12 @@ class ContinuousBatchProcessor:
             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)  # 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)
+            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)
 
 
 # Manager Class (User Interface)
@@ -803,9 +727,9 @@ class ContinuousBatchingManager:
         generation_config: GenerationConfig,
         manual_eviction: bool = False,
         max_queue_size: int = 0,
-        num_q_padding_intervals: int = 0,
-        num_kv_padding_intervals: int = 0,
-        allow_block_sharing: bool = True,
+        num_q_cuda_graphs: int = 0,
+        num_kv_cuda_graphs: int = 0,
+        allow_prefix_sharing: bool = True,
     ) -> None:
         """Initialize the continuous batching manager.
 
@@ -813,99 +737,65 @@ class ContinuousBatchingManager:
             model: The language model for generation
             generation_config: Configuration for generation parameters
             max_queue_size: Maximum size of the request queue (0 = unlimited)
-            num_q_padding_intervals: (optional) Number of intervals used to pad the query dimension
-            num_kv_padding_intervals: (optional) Number of intervals used to pad the keys/values dimension
-            allow_block_sharing: (optional) Whether to allow block sharing if the model has some full attention layers
+            num_q_cuda_graphs: (optional) Number of CUDA graphs to use for the query dimension
+            num_kv_cuda_graphs: (optional) Number of CUDA graphs to use for the keys/values dimension
+            allow_prefix_sharing: (optional) Whether to allow prefix sharing if the model has only full attention layers
         """
-        # Reload paged version of the attention implementation if necessary
         if "paged|" not in model.config._attn_implementation:
-            model.set_attn_implementation(f"paged|{model.config._attn_implementation}")
+            attn_implementation = f"paged|{model.config._attn_implementation}"
+            model.config._attn_implementation = attn_implementation
 
-        # Internal arguments
-        self.model = model.eval()
-        self.manual_eviction = manual_eviction
-        self._allow_block_sharing = allow_block_sharing
-        self._use_prefix_sharing = allow_block_sharing  # approximation until the cache is created
+            # lazy loading flash attention including kernel variations
+            if "flash" in attn_implementation:
+                from ...modeling_flash_attention_utils import lazy_import_paged_flash_attention
 
+                lazy_import_paged_flash_attention(attn_implementation)
+
+        self.model = model.eval()
+        generation_config = model.generation_config if generation_config is None else generation_config
+        self.generation_config = generation_config
         self.input_queue = queue.Queue(maxsize=max_queue_size)
         self.output_queue = queue.Queue()
         self.stop_event = threading.Event()
-        self.batch_processor: ContinuousBatchProcessor | None = None
+        self.log_prob_generation = getattr(generation_config, "log_prob_generation", False)
         self._generation_thread = None
         self._request_counter = 0
         self._request_lock = threading.Lock()
-
-        # Generation config related arguments
-        generation_config = model.generation_config if generation_config is None else generation_config
-        self.generation_config = generation_config
-        self.log_prob_generation = getattr(generation_config, "log_prob_generation", False)
+        self.model.generation_config.top_p = None
         self.do_sample = getattr(generation_config, "do_sample", True)
         self.logit_processor = self.model._get_logits_processor(generation_config)
-        num_return_sequences = getattr(generation_config, "num_return_sequences", None)
-        self.num_return_sequences = num_return_sequences if num_return_sequences is not None else 1
+        use_cuda_graph: bool | None = getattr(generation_config, "use_cuda_graph", None)
+        self.profile = getattr(generation_config, "profile", False)  # TODO: not supported yet
+        self.manual_eviction = manual_eviction
+        self.batch_processor: ContinuousBatchProcessor | None = None
 
-        # self.model.generation_config.top_p = None NOTE: figure out why this was here
+        self._allow_prefix_sharing = allow_prefix_sharing
 
-        # Cuda graph behavior is determined below using either user-specified arguments or heuristics
-        self.use_cuda_graph = self._decide_use_cuda_graphs(
-            use_cuda_graph=getattr(generation_config, "use_cuda_graph", None),
-            num_q_padding_intervals=num_q_padding_intervals,
-            num_kv_padding_intervals=num_kv_padding_intervals,
-            compile_config=getattr(generation_config, "compile_config", None),
-        )
+        # If a number of cuda graphs was specified for either Q or KV, we activate cuda graphs
+        if num_q_cuda_graphs > 0 or num_kv_cuda_graphs > 0:
+            self.use_cuda_graph = True
+        # If use_cuda_graph is specified, we follow the user's choice
+        elif use_cuda_graph is not None:
+            self.use_cuda_graph = use_cuda_graph
+        # If the use of cuda graphs is not specified, we follow the user's choice, otherwise we have a default heuristic
+        else:
+            # Attention implementations where an attention mask is needed suffer a lot more from the padding associated
+            # with cuda graphs, so default is to turn cuda graphs off for those implementations
+            self.use_cuda_graph = not attn_mask_is_needed(self.model.config)
+            logger.warning(
+                f"No behavior specified for use_cuda_graph, defaulting to {self.use_cuda_graph = } because "
+                f"{self.model.config._attn_implementation = }. If you want to save memory, turn off cuda graphs, but "
+                "they can improve performances."
+            )
 
-        # We set the number of padding intervals for Q and KV
-        self.q_padding_intervals = num_q_padding_intervals if num_q_padding_intervals > 0 else NUM_Q_PADDING_INTERVALS
-        self.kv_padding_intervals = (
-            num_kv_padding_intervals if num_kv_padding_intervals > 0 else NUM_KV_PADDING_INTERVALS
-        )
+        # If cuda graphs are activated, we set the number of cuda graphs for Q and KV if not specified
+        if self.use_cuda_graph:
+            self.num_q_cuda_graphs = num_q_cuda_graphs if num_q_cuda_graphs > 0 else NUM_Q_CUDA_GRAPHS
+            self.num_kv_cuda_graphs = num_kv_cuda_graphs if num_kv_cuda_graphs > 0 else NUM_KV_CUDA_GRAPHS
 
-        # Log probability generation is not supported yet (TODO)
         if self.log_prob_generation:
             raise NotImplementedError("log_prob_generation is not supported yet")
 
-    def _decide_use_cuda_graphs(
-        self,
-        use_cuda_graph: bool | None,
-        num_q_padding_intervals: int,
-        num_kv_padding_intervals: int,
-        compile_config: CompileConfig | None,
-    ) -> bool:
-        """Returns whether or not to use cuda graphs for continuous batching, depending on the following criteria:
-        - (use_cuda_graph) which is the user choice
-        - (num_q_padding_intervals) or (num_kv_padding_intervals) which is used to pad inputs: if it was specified by
-            the user, it's probable they want to use cuda graphs so inputs need to be padded
-        - (compile_config): if compile is on, turn on cuda graphs unless the compile mode uses its own cudagraphs
-        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 use_cuda_graph is specified, we follow the user's choice
-        if use_cuda_graph is not None:
-            return use_cuda_graph
-        # If a number of padding intervals was specified for either Q or KV, we activate cuda graphs
-        if num_q_padding_intervals > 0 or num_kv_padding_intervals > 0:
-            return True
-        # If a compile config was found, turn off cuda graphs if the compile config already uses them
-        if compile_config is not None:
-            options = torch._inductor.list_mode_options().get(compile_config.mode, compile_config.options)
-            compile_uses_cudagraphs = options.get("triton.cudagraphs", False)
-            if compile_uses_cudagraphs:
-                logger.warning(
-                    f"Compile config {compile_config.mode = } uses cudagraphs, which usually does not work well with "
-                    "continuous batching. We recommend using mode 'default' or 'max-autotune-no-cudagraphs' instead."
-                )
-            return not compile_uses_cudagraphs  # TODO: should this also match the dynamic shapes?
-        # Otherwise we have a default heuristic based on the attention implementation:
-        # attention implementations where an attention mask is needed suffer a lot more from the padding associated
-        # with cuda graphs, so default is to turn cuda graphs off for those implementations
-        use_cuda_graph = not attn_mask_is_needed(self.model.config)
-        logger.warning(
-            f"No behavior specified for use_cuda_graph, defaulting to {use_cuda_graph = } because "
-            f"{self.model.config._attn_implementation = }. If you want to save memory, turn off cuda graphs, but "
-            "they can improve performances."
-        )
-        return use_cuda_graph
-
     @traced
     def start(self) -> None:
         """Start the background generation thread."""
@@ -932,7 +822,7 @@ class ContinuousBatchingManager:
             logger.warning("\nBatch processor was not initialized.")
         else:
             if self.batch_processor.cache.use_prefix_sharing:
-                logger.info(
+                logger.warning(
                     f"\nPrefix sharing was on. Total prefix length: {self.batch_processor.cache._total_prefix_length}"
                 )
 
@@ -994,7 +884,6 @@ class ContinuousBatchingManager:
         state = RequestState(
             request_id=request_id,
             initial_tokens=list(input_ids),
-            num_children=self.num_return_sequences - 1,
             record_timestamps=record_timestamps,
             tokens_to_process=list(input_ids),
             max_new_tokens=max_new_tokens,
@@ -1013,10 +902,6 @@ class ContinuousBatchingManager:
         streaming: bool = False,
         record_timestamps: bool = False,
     ) -> None:
-        # If there is prefix sharing, we sort the inputs to maximize cache hits
-        if self._use_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
@@ -1087,9 +972,8 @@ class ContinuousBatchingManager:
                 self.model.device,
                 self.model.dtype,
                 tp_size=getattr(self.model, "_tp_size", None),  # Use model's actual TP setting
-                allow_block_sharing=self._allow_block_sharing,
+                allow_prefix_sharing=self._allow_prefix_sharing,
             )
-            self._use_prefix_sharing = paged_attention_cache.use_prefix_sharing  # update the approximation
             logger.debug(f"PagedAttentionCache created in {perf_counter() - t0} seconds")
 
             scheduler = None
@@ -1115,8 +999,6 @@ class ContinuousBatchingManager:
                 scheduler=scheduler(paged_attention_cache, self.manual_eviction),
                 manual_eviction=self.manual_eviction,
                 use_cuda_graph=self.use_cuda_graph,
-                q_padding_intervals=self.q_padding_intervals,
-                kv_padding_intervals=self.kv_padding_intervals,
             )
             self.batch_processor = batch_processor
             self.current_batch = 0
@@ -1139,15 +1021,16 @@ class ContinuousBatchingManager:
         # Loop body ends if there is no requests in the batch
         if not batch_processor.prepare_next_batch():
             return
-        # Debug logging of the current memory usage -- commented out because it's often not used even in debug
-        # if logger.level < logging.DEBUG:
-        #     device, total, reserved, allocated = get_device_and_memory_breakdown()
-        #     available_memory = total - max(allocated, reserved)
-        #     logger.debug(
-        # f"[Memory] Device: {device}, Total: {total}, Reserved: {reserved}, Allocated: {allocated}, Available: {available_memory}"
-        #     )
+        # Debug logging of the current memory usage
+        if logger.level <= logging.DEBUG:
+            device, total, reserved, allocated = get_device_and_memory_breakdown()
+            logger.debug(f"[Memory] Device: {device}, Total: {total}, Reserved: {reserved}, Allocated: {allocated}")
 
         self._generation_step()
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        # Processor updates the batch after generation step is truly over
         batch_processor.update_batch()
 
     @traced
@@ -1189,7 +1072,7 @@ class ContinuousMixin:
         max_queue_size: int = 0,
         num_q_cuda_graphs: int = 0,
         num_kv_cuda_graphs: int = 0,
-        allow_block_sharing: bool = True,
+        allow_prefix_sharing: bool = True,
         block: bool = True,
         timeout: float | None = None,
     ) -> Generator[ContinuousBatchingManager]:
@@ -1199,7 +1082,7 @@ class ContinuousMixin:
             max_queue_size,
             num_q_cuda_graphs,
             num_kv_cuda_graphs,
-            allow_block_sharing,
+            allow_prefix_sharing,
         )
         manager.start()
         try:
@@ -1216,19 +1099,18 @@ class ContinuousMixin:
         generation_config: GenerationConfig | None = None,
         manual_eviction: bool = False,
         max_queue_size: int = 0,
-        num_q_padding_intervals: int = 0,
-        num_kv_padding_intervals: int = 0,
-        allow_block_sharing: bool = True,
+        num_q_cuda_graphs: int = 0,
+        num_kv_cuda_graphs: int = 0,
+        allow_prefix_sharing: bool = True,
     ) -> ContinuousBatchingManager:
         """Initialize a manager for continuous batching inference.
 
         Args:
-            generation_config: An optional generation configuration, which may contain a CompileConfig object
+            generation_config: Custom generation configuration
             manual_eviction: Whether to manually evict requests from the cache
             max_queue_size: Maximum size of the input request queue
-            num_q_padding_intervals: Number of intervals used to pad the query dimension
-            num_kv_padding_intervals: Number of intervals used to pad the keys/values dimension
-            allow_block_sharing: A flag to allow block sharing if the model has some full attention layers
+            num_q_cuda_graphs: Number of CUDA graphs to use for the query dimension
+            num_kv_cuda_graphs: Number of CUDA graphs to use for the keys/values dimension
 
         Returns:
             `ContinuousBatchingManager`: The manager instance to add requests and retrieve results.
@@ -1250,9 +1132,9 @@ class ContinuousMixin:
             generation_config=gen_config,
             manual_eviction=manual_eviction,
             max_queue_size=max_queue_size,
-            num_q_padding_intervals=num_q_padding_intervals,
-            num_kv_padding_intervals=num_kv_padding_intervals,
-            allow_block_sharing=allow_block_sharing,
+            num_q_cuda_graphs=num_q_cuda_graphs,
+            num_kv_cuda_graphs=num_kv_cuda_graphs,
+            allow_prefix_sharing=allow_prefix_sharing,
         )
 
     # TODO: support streaming
@@ -1262,11 +1144,11 @@ class ContinuousMixin:
         self,
         inputs: list[list[int]],
         generation_config: GenerationConfig | None = None,
-        num_q_padding_intervals: int = 0,
-        num_kv_padding_intervals: int = 0,
-        allow_block_sharing: bool = True,
-        record_timestamps: bool = False,
         progress_bar: bool = True,
+        num_q_cuda_graphs: int = 0,
+        num_kv_cuda_graphs: int = 0,
+        allow_prefix_sharing: bool = True,
+        record_timestamps: bool = False,
         **kwargs,
     ) -> dict[str, GenerationOutput]:
         """Generate sequences for a batch of prompts using continuous batching.
@@ -1274,15 +1156,14 @@ class ContinuousMixin:
         Args:
             inputs: List of input token sequences (prompts)
             generation_config: Optional generation configuration
-            num_q_padding_intervals: Number of intervals used to pad the query dimension
-            num_kv_padding_intervals: Number of intervals used to pad the keys/values dimension
-            allow_block_sharing: A flag to allow block sharing if the model has some full attention layers
-            record_timestamps: If set to true, the requests will have a timestamp for each token generated
-            progress_bar: If set to true, a progress bar will be displayed
+            num_q_cuda_graphs: Number of CUDA graphs to use for the query dimension
+            num_kv_cuda_graphs: Number of CUDA graphs to use for the keys/values dimension
             **kwargs: Additional generation parameters
 
         Returns:
-            `dict[str, GenerationOutput]`: a dictionary of request ids to GenerationOutput objects
+            `list[list[int]]`: A list containing the generated sequences (including prompt tokens
+                                if not handled otherwise) for each input prompt, in the same order.
+                                Returns an empty list `[]` for requests that failed.
         """
         if not inputs:
             return {}
@@ -1292,30 +1173,26 @@ class ContinuousMixin:
 
         # Initialize manager with the batch inputs
         results = {}
-        gen_cfg = self.generation_config if generation_config is None else generation_config
-        num_requests = len(inputs) * (gen_cfg.num_return_sequences if gen_cfg.num_return_sequences is not None else 1)
-        # Prepare context managers for the main loop
-        manager_cm = self.continuous_batching_context_manager(
-            generation_config=generation_config,
-            num_q_cuda_graphs=num_q_padding_intervals,
-            num_kv_cuda_graphs=num_kv_padding_intervals,
-            allow_block_sharing=allow_block_sharing,
-            block=True,
-            timeout=5,
-        )
-        logging_cm = logging_redirect_tqdm([logger])
-        pbar_cm = tqdm(
-            total=num_requests,
-            disable=(not progress_bar),
-            desc=f"Solving {num_requests} requests",
-            unit="request",
-        )
-        # Main loop
-        with manager_cm as manager, logging_cm, pbar_cm as pbar:
+        num_requests = len(inputs)
+        with (
+            self.continuous_batching_context_manager(
+                generation_config=generation_config,
+                num_q_cuda_graphs=num_q_cuda_graphs,
+                num_kv_cuda_graphs=num_kv_cuda_graphs,
+                allow_prefix_sharing=allow_prefix_sharing,
+                block=True,
+                timeout=5,
+            ) as manager,
+            logging_redirect_tqdm([logger]),
+            tqdm(
+                total=num_requests,
+                disable=(not progress_bar),
+                desc=f"Solving {num_requests} requests",
+                unit="request",
+            ) as pbar,
+        ):
             try:
-                manager.add_requests(
-                    inputs=inputs, max_new_tokens=kwargs.get("max_new_tokens"), record_timestamps=record_timestamps
-                )
+                manager.add_requests(inputs=inputs, max_new_tokens=kwargs.get("max_new_tokens"))
                 finished_count = 0
                 while finished_count < num_requests:
                     result = manager.get_result(timeout=1)