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

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 14, 05:30 UTC

[Core][KVConnector] Propagate all tokens on resumed preemptions (#24926)

Signed-off-by: Qier Li <kevin44036@gmail.com> Co-authored-by: Qier Li <qier@fb.com>

Qier Li committed Oct 9, 2025 at 02:43 UTC d17f0fbf304fde851e1f4995a30d78182985044e

4 files changed +60 -9

tests/v1/core/test_scheduler.py

+42 -4

@@ -1950,7 +1950,7 @@ def test_schedule_skip_tokenizer_init_structured_output_request():

1950 assert len(scheduler.waiting) == 1

1951

1952

1953 -def test_priority_scheduling_preemption_when_out_of_kv():

1953 +def test_priority_scheduling_preemption_and_resumption_when_out_of_kv():

1954 """Test that priority scheduling preempts lower priority requests

1955 when out of KV cache space."""

1956 # Create scheduler with very limited memory to force preemption

@@ -1959,6 +1959,7 @@ def test_priority_scheduling_preemption_when_out_of_kv():

1959 max_num_batched_tokens=200,

1960 num_blocks=5, # Can hold 64 tokens (first block is null)

1961 block_size=16, # Standard block size

1962 + use_kv_connector=True,

1963 )

1964

1965 # Create a request and schedule it

@@ -1970,12 +1971,13 @@ def test_priority_scheduling_preemption_when_out_of_kv():

1971 starting_idx=0,

1972 )[0]

1973 scheduler.add_request(request_low)

1974 + # 1st schedule

1975 output = scheduler.schedule()

1976 assert len(output.scheduled_new_reqs) == 1

1977 assert len(scheduler.waiting) == 0

1978 assert len(scheduler.running) == 1

1979

1978 - # Simulate model execution

1980 + # Simulate model execution - 1st decode

1981 model_output = ModelRunnerOutput(

1982 req_ids=[request_low.request_id],

1983 req_id_to_index={request_low.request_id: 0},

@@ -1996,6 +1998,7 @@ def test_priority_scheduling_preemption_when_out_of_kv():

1998 starting_idx=1,

1999 )[0]

2000 scheduler.add_request(request_high)

2001 + # 2nd schedule

2002 output = scheduler.schedule()

2003 # KV cache should be full at this point

2004 assert scheduler.kv_cache_manager.block_pool.get_num_free_blocks() == 0

@@ -2004,7 +2007,7 @@ def test_priority_scheduling_preemption_when_out_of_kv():

2007 assert len(scheduler.waiting) == 0

2008 assert len(scheduler.running) == 2

2009

2007 - # Simulate model execution

2010 + # Simulate model execution - 2nd decode

2011 requests = [request_low, request_high]

2012 model_output = ModelRunnerOutput(

2013 req_ids=[req.request_id for req in requests],

@@ -2017,7 +2020,7 @@ def test_priority_scheduling_preemption_when_out_of_kv():

2020 )

2021 scheduler.update_from_output(output, model_output)

2022

2020 - # Schedule again - this should trigger preemption

2023 + # 3rd schedule - this should trigger preemption

2024 # req_low needs 32 tokens = 2 blocks

2025 # req_high needs 33 tokens = 3 blocks

2026 # so doesn't fit in 4 blocks.

@@ -2027,9 +2030,44 @@ def test_priority_scheduling_preemption_when_out_of_kv():

2030 assert len(output.scheduled_new_reqs) == 0

2031 assert output.scheduled_cached_reqs.num_reqs == 1

2032 assert output.scheduled_cached_reqs.req_ids[0] == request_high.request_id

2033 + assert scheduler.requests[request_low.request_id].status == RequestStatus.PREEMPTED

2034 assert len(scheduler.waiting) == 1

2035 assert len(scheduler.running) == 1

2036

2037 + # Simulate model execution - 3rd decode

2038 + model_output = ModelRunnerOutput(

2039 + req_ids=[req.request_id for req in requests],

2040 + req_id_to_index={req.request_id: i for i, req in enumerate(requests)},

2041 + sampled_token_ids=[[], [100]],

2042 + # spec_token_ids=None,

2043 + logprobs=None,

2044 + prompt_logprobs_dict={},

2045 + pooler_output=[],

2046 + )

2047 + # Finish the requests to make room for the preempted requests to resume

2048 + scheduler.update_from_output(output, model_output)

2049 + scheduler.finish_requests(request_high.request_id, RequestStatus.FINISHED_STOPPED)

2050 +

2051 + # 4th Schedule - this should trigger the resumption

2052 + output = scheduler.schedule()

2053 + scheduled_cached_reqs = output.scheduled_cached_reqs

2054 + resumed_from_preemption = scheduled_cached_reqs.resumed_from_preemption

2055 +

2056 + assert len(output.scheduled_new_reqs) == 0

2057 + assert scheduled_cached_reqs.num_reqs == 1

2058 + assert len(scheduler.waiting) == 0

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

2060 +

2061 + # Preempted request resumed in scheduled_cached_reqs

2062 + assert len(resumed_from_preemption) == 1

2063 + assert len(scheduled_cached_reqs.resumed_req_token_ids) == 1

2064 + assert resumed_from_preemption[0]

2065 + assert scheduled_cached_reqs.req_ids[0] == request_low.request_id

2066 + assert scheduled_cached_reqs.resumed_req_token_ids[0] is not None

2067 + # Resumed tokens include 30 prompt tokens and 2 decoded tokens

2068 + assert len(scheduled_cached_reqs.resumed_req_token_ids[0]) == 32

2069 + assert scheduled_cached_reqs.resumed_req_token_ids[0][31] == 100

2070 +

2071

2072 @pytest.mark.parametrize(

2073 ("enable_chunked_prefill", "is_encoder_decoder", "expect_enabled"),

tests/v1/worker/test_gpu_model_runner.py

@@ -257,6 +257,7 @@ def test_update_states_request_resumed(model_runner, dist_init):

257 req_ids=[req_id],

258 resumed_from_preemption=[False],

259 new_token_ids=[[]],

260 + resumed_req_token_ids=[None],

261 new_block_ids=([[0]],),

262 num_computed_tokens=[0],

263 num_output_tokens=[0],

vllm/v1/core/sched/output.py

@@ -98,6 +98,9 @@ class CachedRequestData:

98 # NOTE(woosuk): new_token_ids is only used for pipeline parallelism.

99 # When PP is not used, new_token_ids will be empty.

100 new_token_ids: list[list[int]]

101 + # If resumed_from_preemption is True, propogate the token ids to the

102 + # connector, otherwise will be empty.

103 + resumed_req_token_ids: list[list[int] | None]

104 new_block_ids: list[tuple[list[int], ...] | None]

105 num_computed_tokens: list[int]

106 num_output_tokens: list[int]

@@ -112,6 +115,7 @@ class CachedRequestData:

115 req_ids=[],

116 resumed_from_preemption=[],

117 new_token_ids=[],

118 + resumed_req_token_ids=[],

119 new_block_ids=[],

120 num_computed_tokens=[],

121 num_output_tokens=[],

vllm/v1/core/sched/scheduler.py

+13 -5

@@ -709,10 +709,15 @@ class Scheduler(SchedulerInterface):

709 req_ids: list[str] = []

710 new_token_ids: list[list[int]] = []

711 new_block_ids: list[tuple[list[int], ...] | None] = []

712 + resumed_req_token_ids: list[list[int] | None] = []

713 num_computed_tokens: list[int] = []

714 num_output_tokens: list[int] = []

715

715 - for req in itertools.chain(running_reqs, resumed_reqs):

716 + # Because resumed_reqs is usually empty, it is more efficient to do

717 + # in-place appending so that we don't need to allocate a new list.

718 + resumed_from_preemption = [False] * len(running_reqs)

719 + resumed_from_preemption += [True] * len(resumed_reqs)

720 + for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)):

721 req_id = req.request_id

722 req_ids.append(req_id)

723 num_tokens = num_scheduled_tokens[req_id] - len(

@@ -728,20 +733,23 @@ class Scheduler(SchedulerInterface):

733 req.num_computed_tokens : req.num_computed_tokens + num_tokens

734 ]

735 new_token_ids.append(token_ids)

736 + resumed_token_ids = None

737 + if resumed_from_preemption[idx]:

738 + resumed_token_ids = req.all_token_ids[

739 + : req.num_computed_tokens + num_tokens

740 + ]

741 + resumed_req_token_ids.append(resumed_token_ids)

742 new_block_ids.append(

743 req_to_new_blocks[req_id].get_block_ids(allow_none=True)

744 )

745 num_computed_tokens.append(req.num_computed_tokens)

746 num_output_tokens.append(req.num_output_tokens)

736 - # Because resumed_reqs is usually empty, it is more efficient to do

737 - # in-place appending so that we don't need to allocate a new list.

738 - resumed_from_preemption = [False] * len(running_reqs)

739 - resumed_from_preemption += [True] * len(resumed_reqs)

747

748 return CachedRequestData(

749 req_ids=req_ids,

750 resumed_from_preemption=resumed_from_preemption,

751 new_token_ids=new_token_ids,

752 + resumed_req_token_ids=resumed_req_token_ids,

753 new_block_ids=new_block_ids,

754 num_computed_tokens=num_computed_tokens,

755 num_output_tokens=num_output_tokens,