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

Code before:
            elif state.status == RequestStatus.SPLIT_PENDING_REMAINDER:
                state.status = RequestStatus.PREFILLING_SPLIT
            state.remaining_prefill_tokens = request_tokens[token_budget:]
            state.tokens_to_process = request_tokens[:token_budget]

    def _process_candidates(
        self,
        candidates: list[RequestState],
        token_budget: int,
        cache_budget: int,
        request_ids_to_remove_from_waiting: set[str],
        safety_margin: float | None = None,
    ) -> tuple[list[RequestState], bool]:
        """Schedules candidate requests for the current batch.

        This method contains the common logic shared by all schedulers: it checks token and cache budgets, allocates
        cache blocks if needed, updates request states, and tracks which waiting requests should be removed from the
        waiting queue.
        """
        scheduled_requests = []
        one_allocation_failed = False
        safety_margins = safety_margin * self.cache.num_blocks if safety_margin is not None else None

        for state in candidates:
            num_free_blocks = self.cache.get_num_free_blocks()
            if safety_margins is not None:
                # If we are out the safety margin, we only accept decoding requests or the first prefill request
                outside_safety_margin = num_free_blocks < safety_margins
                if outside_safety_margin and scheduled_requests and state.status != RequestStatus.DECODING:
                    logger.info(
                        f"Outside safety margin, breaking out of scheduling loop. {num_free_blocks = } {safety_margins = }"
                    )
                    break

            # Check cache budget
            cache_needed = state.current_len()
            cache_needed = (
                cache_needed if self.cache_budget_module is None else cache_needed % self.cache_budget_module
            )
            if cache_budget < cache_needed:
                continue

            # Infer the tokens that will be present in the batch if token budget is enough
            request_tokens = self._infer_request_tokens(state, request_ids_to_remove_from_waiting)
            # Account for token budget
            request_len = min(len(request_tokens), token_budget)
            # Check there will be enough cache for the new tokens
            allocation_successful = self._allocate_blocks_if_needed(state, request_len)

            # If the allocation would not be successful, we move on to the next request
            if not allocation_successful:

Review comment (suggestion):
remove the check because `safety_margins` will never be None

Code after:
            elif state.status == RequestStatus.SPLIT_PENDING_REMAINDER:
                state.status = RequestStatus.PREFILLING_SPLIT
            state.remaining_prefill_tokens = request_tokens[token_budget:]
            state.tokens_to_process = request_tokens[:token_budget]

    def _process_candidates(
        self,
        candidates: list[RequestState],
        token_budget: int,
        cache_budget: int,
        request_ids_to_remove_from_waiting: set[str],
        safety_margin: float = 0.0,
    ) -> tuple[list[RequestState], bool]:
        """Schedules candidate requests for the current batch.

        This method contains the common logic shared by all schedulers: it checks token and cache budgets, allocates
        cache blocks if needed, updates request states, and tracks which waiting requests should be removed from the
        waiting queue.
        """
        scheduled_requests = []
        one_allocation_failed = False
        safety_margins = safety_margin * self.cache.num_blocks

        for state in candidates:
            num_free_blocks = self.cache.get_num_free_blocks()
            # If we are out the safety margin, we only accept decoding requests or the first prefill request
            outside_safety_margin = num_free_blocks < safety_margins
            if outside_safety_margin and scheduled_requests and state.status != RequestStatus.DECODING:
                logger.info(
                    f"Outside safety margin, breaking out of scheduling loop. {num_free_blocks = } {safety_margins = }"
                )
                break

            # Check cache budget
            cache_needed = state.current_len()
            cache_needed = (
                cache_needed if self.cache_budget_module is None else cache_needed % self.cache_budget_module
            )
            if cache_budget < cache_needed:
                continue

            # Infer the tokens that will be present in the batch if token budget is enough
            request_tokens = self._infer_request_tokens(state, request_ids_to_remove_from_waiting)
            # Account for token budget
            request_len = min(len(request_tokens), token_budget)
            # Check there will be enough cache for the new tokens
            allocation_successful = self._allocate_blocks_if_needed(state, request_len)

            # If the allocation would not be successful, we move on to the next request
            if not allocation_successful:
                one_allocation_failed = True