@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f45a332886

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 02:15 UTC

[Sched] Enhance the logic to remove stopped requests from queues (#20739)

Woosuk Kwon committed Jul 12, 2025 at 15:33 UTC f45a33288624f30e3ab3d1f5c3c4499e86c5bb7d

3 files changed +92 -17

requirements/common.txt

+1 -1

@@ -7,7 +7,7 @@ requests >= 2.26.0

7 tqdm

8 blake3

9 py-cpuinfo

10 -transformers >= 4.51.1

10 +transformers >= 4.53.2

11 huggingface-hub[hf_xet] >= 0.33.0 # Required for Xet downloads.

12 tokenizers >= 0.21.1 # Required for fast incremental detokenization.

13 protobuf # Required by LlamaTokenizer.

tests/v1/core/test_scheduler.py

+62

@@ -451,6 +451,7 @@ def test_stop_via_update_from_output():

451 req.num_computed_tokens = req.num_tokens

452 scheduler.requests[req.request_id] = req

453 scheduler.running.append(req)

454 + req.status = RequestStatus.RUNNING

455

456 scheduler_output = SchedulerOutput(

457 scheduled_new_reqs=[],

@@ -504,6 +505,7 @@ def test_stop_via_update_from_output():

505 req.num_computed_tokens = req.num_tokens

506 scheduler.requests[req.request_id] = req

507 scheduler.running.append(req)

508 + req.status = RequestStatus.RUNNING

509

510 scheduler_output = SchedulerOutput(

511 scheduled_new_reqs=[],

@@ -556,6 +558,7 @@ def test_stop_via_update_from_output():

558 req.num_computed_tokens = req.num_tokens

559 scheduler.requests[req.request_id] = req

560 scheduler.running.append(req)

561 + req.status = RequestStatus.RUNNING

562

563 scheduler_output = SchedulerOutput(

564 scheduled_new_reqs=[],

@@ -703,6 +706,65 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],

706 scheduler.update_from_output(scheduler_output1, model_runner_output)

707

708

709 +def test_preempt_during_execution():

710 + # NOTE(woosuk): The actual number of available blocks is 10 instead of 11

711 + # because block 0 is reserved as the null block.

712 + scheduler = create_scheduler(max_num_batched_tokens=100,

713 + block_size=16,

714 + num_blocks=11,

715 + enable_prefix_caching=False)

716 + requests = create_requests(num_requests=2, num_tokens=80)

717 +

718 + # Schedule the first request.

719 + scheduler.add_request(requests[0])

720 + scheduler_output0 = scheduler.schedule()

721 + assert len(scheduler_output0.num_scheduled_tokens) == 1

722 + assert len(scheduler_output0.scheduled_new_reqs[0].block_ids[0]) == 5

723 +

724 + # Schedule the second request while the first request is still running.

725 + # This scenario can occur in certain cases, when max_concurrent_batches > 1

726 + # (e.g., when pipeline parallelism is used).

727 + scheduler.add_request(requests[1])

728 + scheduler_output1 = scheduler.schedule()

729 + assert len(scheduler_output1.num_scheduled_tokens) == 1

730 + assert len(scheduler_output1.scheduled_new_reqs[0].block_ids[0]) == 5

731 +

732 + # Get the output of the first request.

733 + model_runner_output0 = ModelRunnerOutput(

734 + req_ids=[requests[0].request_id],

735 + req_id_to_index={requests[0].request_id: 0},

736 + sampled_token_ids=[[0]],

737 + spec_token_ids=None,

738 + logprobs=None,

739 + prompt_logprobs_dict={},

740 + pooler_output=[],

741 + )

742 + scheduler.update_from_output(scheduler_output0, model_runner_output0)

743 +

744 + # Schedule the first request again. This will cause the preemption

745 + # of the second request because the KV cache is full.

746 + _ = scheduler.schedule()

747 + assert len(scheduler.running) == 1

748 + assert scheduler.running[0] == requests[0]

749 + assert requests[1].status == RequestStatus.PREEMPTED

750 +

751 + model_runner_output1 = ModelRunnerOutput(

752 + req_ids=[requests[1].request_id],

753 + req_id_to_index={requests[1].request_id: 0},

754 + sampled_token_ids=[[42]],

755 + spec_token_ids=None,

756 + logprobs=None,

757 + prompt_logprobs_dict={},

758 + pooler_output=[],

759 + )

760 + scheduler.update_from_output(scheduler_output1, model_runner_output1)

761 +

762 + # The second request (that is preempted) should be updated with the

763 + # sampled token id.

764 + assert len(requests[1].output_token_ids) == 1

765 + assert requests[1].output_token_ids[0] == 42

766 +

767 +

768 # Note - these test cases mirror some of those in test_rejection_sampler.py

769 @pytest.mark.parametrize(

770 "spec_tokens,output_tokens,expected",

vllm/v1/core/sched/scheduler.py

+29 -16

@@ -747,19 +747,21 @@ class Scheduler(SchedulerInterface):

747 pooler_outputs = model_runner_output.pooler_output

748 num_nans_in_logits = model_runner_output.num_nans_in_logits

749

750 - new_running: list[Request] = []

750 outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)

751 spec_decoding_stats: Optional[SpecDecodingStats] = None

752

754 - # NOTE(woosuk): As len(self.running) can be up to 1K or more, the below

755 - # loop can be a performance bottleneck. We should do our best to avoid

756 - # expensive operations inside the loop.

757 - for request in self.running:

758 - req_id = request.request_id

759 - num_tokens_scheduled = num_scheduled_tokens.get(req_id, 0)

760 - if num_tokens_scheduled == 0:

761 - # The request was not scheduled in this step.

762 - new_running.append(request)

753 + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,

754 + # the below loop can be a performance bottleneck. We should do our best

755 + # to avoid expensive operations inside the loop.

756 + stopped_running_reqs: set[Request] = set()

757 + stopped_preempted_reqs: set[Request] = set()

758 + for req_id, num_tokens_scheduled in num_scheduled_tokens.items():

759 + assert num_tokens_scheduled > 0

760 + request = self.requests.get(req_id)

761 + if request is None:

762 + # The request is already finished. This can happen if the

763 + # request is aborted while the model is executing it (e.g.,

764 + # in pipeline parallelism).

765 continue

766

767 req_index = model_runner_output.req_id_to_index[req_id]

@@ -792,6 +794,7 @@ class Scheduler(SchedulerInterface):

794 new_logprobs = None

795 new_token_ids = generated_token_ids

796 kv_transfer_params = None

797 + status_before_stop = request.status

798

799 # Append generated tokens and check for stop. Note that if

800 # a request is still being prefilled, we expect the model runner

@@ -803,17 +806,22 @@ class Scheduler(SchedulerInterface):

806 # This must be called before we make the EngineCoreOutput.

807 stopped = check_stop(request, self.max_model_len)

808 if stopped:

806 - kv_transfer_params = self._free_request(request)

809 del new_token_ids[num_new:] # Trim new tokens if needed.

810 break

811

812 + # Stop checking for pooler models.

813 pooler_output = None

814 if pooler_outputs:

815 pooler_output = pooler_outputs[req_index]

816 stopped = check_stop(request, self.max_model_len,

817 pooler_output)

815 - if stopped:

816 - kv_transfer_params = self._free_request(request)

818 +

819 + if stopped:

820 + kv_transfer_params = self._free_request(request)

821 + if status_before_stop == RequestStatus.RUNNING:

822 + stopped_running_reqs.add(request)

823 + else:

824 + stopped_preempted_reqs.add(request)

825

826 # Extract sample logprobs if needed.

827 if request.sampling_params is not None \

@@ -868,9 +876,14 @@ class Scheduler(SchedulerInterface):

876 # Invariant: EngineCore returns no partial prefill outputs.

877 assert not prompt_logprobs_tensors

878

871 - if not stopped:

872 - new_running.append(request)

873 - self.running = new_running

879 + # Remove the stopped requests from the running and waiting queues.

880 + if stopped_running_reqs:

881 + self.running = [

882 + req for req in self.running if req not in stopped_running_reqs

883 + ]

884 + if stopped_preempted_reqs:

885 + # This is a rare case and unlikely to impact performance.

886 + self.waiting.remove_requests(stopped_preempted_reqs)

887

888 # KV Connector: update state for finished KV Transfers.

889 self._update_from_kv_xfer_finished(model_runner_output)