Repository: huggingface/transformers
File: src/transformers/generation/continuous_batching/continuous_api.py (Python)

Code before:

        # 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 self._graphs is None or not torch.cuda.is_available():
            self._forward_process_and_sample(model, batch_data, logit_processor, do_sample)
            return None

        # If we have a graph that fits, we replay it
        graph = self._graphs.get((padded_q, padded_read_index_size))
        if graph is not None:
            graph.replay()
            return None

        # Otherwise, we need to create it
        logger.info(f"Creating graph for {(padded_q, padded_read_index_size) = }")
        stream = torch.cuda.Stream(device=model.device)
        stream.wait_stream(torch.cuda.current_stream())
        # Warmup
        with torch.cuda.stream(stream):
            self._forward_process_and_sample(model, batch_data, logit_processor, do_sample)
        torch.cuda.current_stream().wait_stream(stream)
        # Catpure
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph, stream=stream):
            self._forward_process_and_sample(model, batch_data, logit_processor, do_sample)
        self._graphs[(padded_q, padded_read_index_size)] = graph

    @traced
    def _forward_process_and_sample(

Review comment (suggestion):
This piece of code is called at every iteration, so it's better not to add calls like this one. Why not change the behavior of `use_cuda_graphs` instead? There is code to decide whether or not to use them, and `torch.cuda.is_available` could be a first condition there.

Code after:

        # 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)

    @traced
    def has_pending_requests(self) -> bool:
        """Check if there are any active or waiting requests."""
        return self.scheduler.has_pending_requests()

    @traced
    def handle_batch_error(self, error):
        """Handle errors during batch processing."""
        failed_reqs = self.requests_in_batch
        for req in failed_reqs:
            self._handle_request_error(error, req)
            self.scheduler.finish_request(req.request_id)

    @traced
    def fail_all_requests(self, error: Exception) -> None:
        """Fail all active requests with the given error.

        Args:
            error: The error to report in the failure message
        """

        requests = list(self.scheduler.active_requests.values())
        for state in requests:
            self._handle_request_error(error, state)
            self.scheduler.finish_request(state.request_id)

        # Also fail any requests in the waiting queue
        for req_id in list(self.scheduler.waiting_requests.keys()):
            state = self.scheduler.waiting_requests.pop(req_id)
            self._handle_request_error(error, state)

        # Clear the ordering queue
        self.scheduler.waiting_requests_order.clear()

    @traced
    @torch.no_grad()
    def _generation_step(self, model: nn.Module, logit_processor: LogitsProcessorList, 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,